From 4dae43a89768ac729012cfee919e826306581fe8 Mon Sep 17 00:00:00 2001 From: Nicholson Koukpaizan Date: Sat, 17 Oct 2026 13:00:13 -0400 Subject: [PATCH] Fix Enzyme sparse_store dispatch and drop the unused primal buffer. --- .../AutomaticDifferentiation/Enzyme/DfDwb.hpp | 24 +++++++------------ .../AutomaticDifferentiation/Enzyme/DfDws.hpp | 12 ++++------ .../AutomaticDifferentiation/Enzyme/DfDy.hpp | 24 +++++++------------ .../AutomaticDifferentiation/Enzyme/DfDyp.hpp | 24 +++++++------------ .../AutomaticDifferentiation/Enzyme/DhDwb.hpp | 12 ++++------ .../AutomaticDifferentiation/Enzyme/DhDy.hpp | 12 ++++------ .../Enzyme/LowerSparseStorage.hpp | 14 +++++------ 7 files changed, 43 insertions(+), 79 deletions(-) diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp index 79e162dcc5..73fa06ea74 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp @@ -59,7 +59,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -77,11 +76,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -91,10 +87,10 @@ namespace GridKit yp, enzyme_dup, wb, - output, + output, // Seed: column var_i of the identity enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } @@ -130,7 +126,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -148,11 +143,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -162,12 +154,12 @@ namespace GridKit yp, enzyme_dup, wb, - output, + output, // Seed: column var_i of the identity enzyme_const, ws, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp index 2439601b19..a3f4cb0345 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp @@ -61,7 +61,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -79,11 +78,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -95,10 +91,10 @@ namespace GridKit wb, enzyme_dup, ws, - output, + output, // Seed: column var_i of the identity enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DfDy.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DfDy.hpp index 21f1060d65..19051d69c0 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DfDy.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DfDy.hpp @@ -59,7 +59,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -77,24 +76,21 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, enzyme_dup, y, - output, + output, // Seed: column var_i of the identity enzyme_const, yp, enzyme_const, wb, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } @@ -130,7 +126,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -148,17 +143,14 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, enzyme_dup, y, - output, + output, // Seed: column var_i of the identity enzyme_const, yp, enzyme_const, @@ -166,8 +158,8 @@ namespace GridKit enzyme_const, ws, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DfDyp.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DfDyp.hpp index f24a250b9f..b7505800fe 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DfDyp.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DfDyp.hpp @@ -61,7 +61,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -79,11 +78,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -91,12 +87,12 @@ namespace GridKit y, enzyme_dup, yp, - output, + output, // Seed: column var_i of the identity enzyme_const, wb, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } @@ -134,7 +130,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -152,11 +147,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -164,14 +156,14 @@ namespace GridKit y, enzyme_dup, yp, - output, + output, // Seed: column var_i of the identity enzyme_const, wb, enzyme_const, ws, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DhDwb.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DhDwb.hpp index ac06194015..93d5048e75 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DhDwb.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DhDwb.hpp @@ -59,7 +59,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -77,11 +76,8 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, @@ -91,10 +87,10 @@ namespace GridKit yp, enzyme_dup, wb, - output, + output, // Seed: column var_i of the identity enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DhDy.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DhDy.hpp index 30cdb98af1..da7cd2b876 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/DhDy.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/DhDy.hpp @@ -59,7 +59,6 @@ namespace GridKit { if (n_res > 0 && n_var > 0) { - std::vector elementary_v(n_var); for (size_t var_i = 0; var_i < n_var; ++var_i) { // Sparse storage. @see LowerSparseStorage.hpp @@ -77,24 +76,21 @@ namespace GridKit vals, &nnz); - // Elementary vector for Jacobian-vector product - std::ranges::fill(elementary_v, 0.0); - elementary_v[var_i] = 1.0; - // Core automatic differentiation intrinsic that will be replaced by a derivative + // The dummy primal pointer is safe only while residuals never read back their output __enzyme_fwddiff((void*) ModelWrapper::eval, enzyme_const, model, enzyme_dup, y, - output, + output, // Seed: column var_i of the identity enzyme_const, yp, enzyme_const, wb, enzyme_dupnoneed, - elementary_v.data(), - d_output); + (ScalarT*) 0x1, // Unneeded primal residual (dummy pointer) + d_output); // Column var_i of the Jacobian } } } diff --git a/GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp b/GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp index 6714a28a77..0c379be266 100644 --- a/GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp +++ b/GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp @@ -216,7 +216,7 @@ namespace GridKit IdxT* rows, IdxT* cols, ScalarT* vals, - IdxT& nnz) + IdxT* nnz) { if (val == 0.0) return; @@ -226,17 +226,17 @@ namespace GridKit // this template nightmare is because __attribute__((enzyme_sparse_accumulate)) does not support templates yet if constexpr (std::is_same::value) { - if constexpr (std::is_same::value) - inner_store_float_size_t(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, nnz); + if constexpr (std::is_same::value) + inner_store_float_size_t(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, *nnz); else - inner_store_double_size_t(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, nnz); + inner_store_double_size_t(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, *nnz); } else if constexpr (std::is_same::value) { - if constexpr (std::is_same::value) - inner_store_float_long_int(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, nnz); + if constexpr (std::is_same::value) + inner_store_float_long_int(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, *nnz); else - inner_store_double_long_int(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, nnz); + inner_store_double_long_int(row, col, val, scaling, row_indices, col_indices, rows, cols, vals, *nnz); } else {