diff --git a/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp b/GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp index 79e162dcc..73fa06ea7 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 2439601b1..a3f4cb034 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 21f1060d6..19051d69c 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 f24a250b9..b7505800f 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 ac0619401..93d5048e7 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 30cdb98af..da7cd2b87 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 6714a28a7..0c379be26 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 {