Skip to content

Nuke COO_Matrix class #447

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
nkoukpaizan wants to merge 39 commits into
base: develop
Choose a base branch
from
Draft

Nuke COO_Matrix class #447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
5df4a2f
WIP: nuke COO_Matrix class. Remove what's not needed and cleanup.
nkoukpaizan Jun 8, 2026
04cda17
Separate implementation and calls for DfDy and DfDyp.
nkoukpaizan Jun 8, 2026
6caca2c
Update Component class templates.
nkoukpaizan Jun 8, 2026
31623ad
Get ScalarT and IdxT from ModelT::ScalarT and ModelT::IdxT in Enzyme …
nkoukpaizan Jun 8, 2026
2cbd8e6
Return nnz from partial derivative Enzyme wrappers.
nkoukpaizan Jun 9, 2026
108dd16
Move csr_jac_ member variable from SystemModel to Component.
nkoukpaizan Jun 9, 2026
bd9e951
Remove sorted_ member variable.
nkoukpaizan Jun 9, 2026
3479621
Store bus Jacobian contributions within Branch, Genrou and Gensal (in…
nkoukpaizan Jun 9, 2026
3acf58d
Make BusFault store currents as internal variables --> can now also s…
nkoukpaizan Jun 9, 2026
106633c
Nuke COO_Matrix::axpy.
nkoukpaizan Jun 9, 2026
1f42515
Bug fix in BusFault: always evaluate the residuals, but only contribu…
nkoukpaizan Jun 9, 2026
ef11cc8
Nuke COO_Matrix::zeroValuedMatrix.
nkoukpaizan Jun 9, 2026
4401cc7
nnz-->block_nnz in component Jacobians.
nkoukpaizan Jun 9, 2026
d7abb5f
Minor edits.
nkoukpaizan Jun 10, 2026
dd137a7
.size()-->nnz_ in COO_Matrix.
nkoukpaizan Jun 10, 2026
c780799
Compute nnz_dup in component models (outside COO_Matrix).
nkoukpaizan Jun 10, 2026
78e929c
PhasorDynamics::SystemModel should not need PhasorDynamics:: to acces…
nkoukpaizan Jun 10, 2026
7f78513
Proof of concept to construct CsrMatrix at component level.
nkoukpaizan Jun 10, 2026
6f68059
constructCsrFromCOO in Component base class.
nkoukpaizan Jun 10, 2026
08c4007
Construct CsrMatrix for all components and update tests. Currently sl…
nkoukpaizan Jun 11, 2026
af468a2
Nuke MapFromCOO.
nkoukpaizan Jun 11, 2026
a09de14
Assemble system Jacobian from component Csr (instead of COO).
nkoukpaizan Jun 11, 2026
abda163
Nuke getJacobian method that returns COO_Matrix.
nkoukpaizan Jun 11, 2026
911e7b3
Only construct component-level CsrMatrix once.
nkoukpaizan Jun 11, 2026
d6b594a
No need for csr_jac_ counter at component level.
nkoukpaizan Jun 11, 2026
9852286
Use alias for CooMatrixT in SystemModel.
nkoukpaizan Jun 11, 2026
8a77560
Remve Jacobian debug print from System model.
nkoukpaizan Jun 11, 2026
71d0e5f
Inline CooMatrix and CsrMatrix size accessors.
nkoukpaizan Jun 11, 2026
ce79f73
Temporary variable for n_rows.
nkoukpaizan Jun 11, 2026
5e03014
A few more inline.
nkoukpaizan Jun 11, 2026
b87cb8e
Include Definitions.hpp in ExciterSexsPtiTests.
nkoukpaizan Jun 11, 2026
3199d50
Move alpha* down to LowerSparseStorage.
nkoukpaizan Jun 11, 2026
8df04da
Fix memory leak.
nkoukpaizan Jun 11, 2026
d4c07d1
Using CooMatrix is faster than CsrMatrix at the component level. Nuke…
nkoukpaizan Jun 11, 2026
9f1a9e3
Apply pre-commit fixes
nkoukpaizan Jun 11, 2026
963fc96
More precise bounds on Jacobian buffers.
nkoukpaizan Jun 11, 2026
dc1ab07
Add a few comments.
nkoukpaizan Jun 11, 2026
6940d61
Revert inlining changes to CooMatrix and CsrMatrix.
nkoukpaizan Jun 11, 2026
cb48731
Fix formatting.
nkoukpaizan Jun 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/pre_commit.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,4 +19,4 @@ jobs:
fetch: false
committer_name: GitHub Actions
committer_email: actions@github.com
message: Apply pre-commmit fixes
message: Apply pre-commit fixes
37 changes: 14 additions & 23 deletions GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@
#include <GridKit/AutomaticDifferentiation/Enzyme/EnzymeDefinitions.hpp>
#include <GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp>
#include <GridKit/AutomaticDifferentiation/Enzyme/ModelWrappers.hpp>
#include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
#include <GridKit/ScalarTraits.hpp>

namespace GridKit
{
Expand All @@ -23,14 +21,13 @@ namespace GridKit
*
* @tparam ModelT - model type
* @tparam MemberFunctions - member function parameter key
* @tparam ScalarT - scalar data type
* @tparam IdxT - matrix index data type
*/
template <typename ModelT, MemberFunctions function, class ScalarT, typename IdxT>
template <typename ModelT, MemberFunctions function>
struct DfDwb
{
using RealT = typename GridKit::ScalarTraits<ScalarT>::RealT;
using MatrixT = GridKit::LinearAlgebra::COO_Matrix<RealT, IdxT>;
using ScalarT = typename ModelT::ScalarT;
using IdxT = typename ModelT::IdxT;
using RealT = typename ModelT::RealT;

/**
* @param[in] model - Pointer to the model to be differentiated
Expand All @@ -41,7 +38,7 @@ namespace GridKit
* @param[in] y - Internal variables
* @param[in] yp - Internal variable derivatives
* @param[in] wb - Bus variables
* @param[in,out] jac - Jacobian
* @param[out] nnz - Number of nonzeros
*/
static void eval(ModelT* model,
size_t n_res,
Expand All @@ -54,12 +51,11 @@ namespace GridKit
IdxT* rows,
IdxT* cols,
RealT* vals,
MatrixT& jac)
IdxT& nnz)
{
if (n_res > 0 && n_var > 0)
{
std::vector<ScalarT> elementary_v(n_var);
IdxT nnz = 0;
for (size_t var_i = 0; var_i < n_var; ++var_i)
{
// Sparse storage. @see LowerSparseStorage.hpp
Expand All @@ -69,6 +65,7 @@ namespace GridKit
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT, IdxT>,
(void*) sparse_store<ScalarT, IdxT>,
var_i,
1.0, // value scaling
res_indices,
var_indices,
rows,
Expand All @@ -80,8 +77,8 @@ namespace GridKit
std::ranges::fill(elementary_v, 0.0);
elementary_v[var_i] = 1.0;

// Core automatic differentiaation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
// Core automatic differentiation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function>::eval,
enzyme_const,
model,
enzyme_const,
Expand All @@ -95,9 +92,6 @@ namespace GridKit
elementary_v.data(),
d_output);
}

// Store result
jac.setValues(1.0, rows, cols, vals, nnz); //< @todo: Update once sparse storage format changes
}
}

Expand All @@ -111,7 +105,7 @@ namespace GridKit
* @param[in] yp - Internal variable derivatives
* @param[in] wb - Bus variables
* @param[in] ws - Signal variables
* @param[in,out] jac - Jacobian
* @param[out] nnz - Number of nonzeros
*/
static void eval(ModelT* model,
size_t n_res,
Expand All @@ -125,12 +119,11 @@ namespace GridKit
IdxT* rows,
IdxT* cols,
RealT* vals,
MatrixT& jac)
IdxT& nnz)
{
if (n_res > 0 && n_var > 0)
{
std::vector<ScalarT> elementary_v(n_var);
IdxT nnz = 0;
for (size_t var_i = 0; var_i < n_var; ++var_i)
{
// Sparse storage. @see LowerSparseStorage.hpp
Expand All @@ -140,6 +133,7 @@ namespace GridKit
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT, IdxT>,
(void*) sparse_store<ScalarT, IdxT>,
var_i,
1.0, // value scaling
res_indices,
var_indices,
rows,
Expand All @@ -151,8 +145,8 @@ namespace GridKit
std::ranges::fill(elementary_v, 0.0);
elementary_v[var_i] = 1.0;

// Core automatic differentiaation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
// Core automatic differentiation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function>::eval,
enzyme_const,
model,
enzyme_const,
Expand All @@ -168,9 +162,6 @@ namespace GridKit
elementary_v.data(),
d_output);
}

// Store result
jac.setValues(1.0, rows, cols, vals, nnz); //< @todo: Update once sparse storage format changes
}
}
};
Expand Down
24 changes: 9 additions & 15 deletions GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@
#include <GridKit/AutomaticDifferentiation/Enzyme/EnzymeDefinitions.hpp>
#include <GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp>
#include <GridKit/AutomaticDifferentiation/Enzyme/ModelWrappers.hpp>
#include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
#include <GridKit/ScalarTraits.hpp>

namespace GridKit
{
Expand All @@ -23,14 +21,13 @@ namespace GridKit
*
* @tparam ModelT - model type
* @tparam MemberFunctions - member function parameter key
* @tparam ScalarT - scalar data type
* @tparam IdxT - matrix index data type
*/
template <typename ModelT, MemberFunctions function, class ScalarT, typename IdxT>
template <typename ModelT, MemberFunctions function>
struct DfDws
{
using RealT = typename GridKit::ScalarTraits<ScalarT>::RealT;
using MatrixT = GridKit::LinearAlgebra::COO_Matrix<RealT, IdxT>;
using ScalarT = typename ModelT::ScalarT;
using IdxT = typename ModelT::IdxT;
using RealT = typename ModelT::RealT;

/**
* @param[in] model - Pointer to the model to be differentiated
Expand All @@ -42,7 +39,7 @@ namespace GridKit
* @param[in] yp - Internal variable derivatives
* @param[in] wb - Bus variables
* @param[in] ws - Signal variables
* @param[in,out] jac - Jacobian
* @param[out] nnz - Number of nonzeros
*/
static void eval(ModelT* model,
size_t n_res,
Expand All @@ -56,12 +53,11 @@ namespace GridKit
IdxT* rows,
IdxT* cols,
RealT* vals,
MatrixT& jac)
IdxT& nnz)
{
if (n_res > 0 && n_var > 0)
{
std::vector<ScalarT> elementary_v(n_var);
IdxT nnz = 0;
for (size_t var_i = 0; var_i < n_var; ++var_i)
{
// Sparse storage. @see LowerSparseStorage.hpp
Expand All @@ -71,6 +67,7 @@ namespace GridKit
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT, IdxT>,
(void*) sparse_store<ScalarT, IdxT>,
var_i,
1.0, // value scaling
res_indices,
var_indices,
rows,
Expand All @@ -82,8 +79,8 @@ namespace GridKit
std::ranges::fill(elementary_v, 0.0);
elementary_v[var_i] = 1.0;

// Core automatic differentiaation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
// Core automatic differentiation intrinsic that will be replaced by a derivative
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function>::eval,
enzyme_const,
model,
enzyme_const,
Expand All @@ -99,9 +96,6 @@ namespace GridKit
elementary_v.data(),
d_output);
}

// Store result
jac.setValues(1.0, rows, cols, vals, nnz); //< @todo: Update once sparse storage format changes
}
}
};
Expand Down
Loading
Loading