Kokkos Automatic Differentiation (Undisplaced) by NamjaeChoi · Pull Request #32597 · idaholab/moose

NamjaeChoi · 2026-04-01T22:35:37Z

Refs #32699.

moosebuild · 2026-04-02T21:38:54Z

Job Documentation, step Docs: sync website on a2f3df9 wanted to post the following:

View the site here

This comment will be updated on new commits.

NamjaeChoi · 2026-05-11T21:24:38Z

They can also do copy construction. It will be just a bit of waste

lindsayad · 2026-05-12T17:14:40Z

    _dot = tag == problem.getVectorTagID(Moose::SOLUTION_DOT_TAG);
+
+  if (problem.vectorTagExists(Moose::OLD_SOLUTION_TAG))
+    _old = _old || tag == problem.getVectorTagID(Moose::OLD_SOLUTION_TAG);


Why the || conditional here?

lindsayad · 2026-05-12T17:15:45Z

    _dot = tag == problem.getVectorTagID(Moose::SOLUTION_DOT_TAG);
+
+  if (problem.vectorTagExists(Moose::OLD_SOLUTION_TAG))
+    _old = _old || tag == problem.getVectorTagID(Moose::OLD_SOLUTION_TAG);


same question

lindsayad · 2026-05-13T16:05:58Z

could we maybe build some unit tests to bring down the 140 unused lines introduced in Coupleable?

lindsayad · 2026-05-13T16:04:19Z

+
+Kokkos-MOOSE kernels also support [automatic differentiation (AD)](automatic_differentiation/index.md).
+AD kernels can be derived from `Moose::Kokkos::ADKernel` and should be registered with `registerKokkosADResidualObject()`.
+AD kernels requires `computeQpResidual()` to be defined with the following signature, where everything remains the same with the ordinary kernels except the return type being `Moose::Kokkos::ADReal`:


Suggested change

AD kernels requires `computeQpResidual()` to be defined with the following signature, where everything remains the same with the ordinary kernels except the return type being `Moose::Kokkos::ADReal`:

AD kernels require `computeQpResidual()` to be defined with the following signature, where everything remains the same with the ordinary kernels except the return type being `Moose::Kokkos::ADReal`:

lindsayad · 2026-05-13T16:04:39Z

+                                                        AssemblyDatum & datum) const;
+```
+
+`computeQpJacobian()` and `computeQpOffDiagJacobian()` are unused, as AD automatically assembles Jacobian.


Suggested change

`computeQpJacobian()` and `computeQpOffDiagJacobian()` are unused, as AD automatically assembles Jacobian.

`computeQpJacobian()` and `computeQpOffDiagJacobian()` are unused, as AD automatically assembles the Jacobian.

lindsayad · 2026-05-13T16:09:02Z

  _dims.copyToDevice();
  _offsets.copyToDevice();
-  _data.create(_offsets.last() + stride);
+  _data.create(_offsets.last() + stride + 1);


was this a bug before? Some comment in the code regarding the +1 here would be nice

lindsayad · 2026-05-13T16:13:47Z

+  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
+  auto row = sys.getElemLocalDofIndex(datum.elem().id, i, _kokkos_var.var(comp));
+
+  for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)


I'd really like to get KOKKOS_FUNCTION support for make_range and index_range (FYI @roystgnr). Because strictly speaking we could be promoting the integer type returned by _matrix_tags.size(), which is the safe way to do it compared to the other other way around. But the pedantic/perfectionist part of me doesn't love it

Wow, now that you mention it, so do I!

I don't want to put it in libMesh/libmesh#4451 (because that's almost ready to get merged, because I've already hassled @rochi00 enough about it, and because it is a logically separate feature) ... although, now that I look at things, I'm quite surprised that Rochi didn't have to get it in there. Looks like, although we #include <int_range.h> in type_vector.h, we only got around to using it for l1_norm() (which his code doesn't use) and we never switched any of the older methods to it!

lindsayad · 2026-05-13T16:21:52Z

+  {
+    auto tag = _matrix_tags[t];
+
+    if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))


So this is to avoid an intermediate matrix close between the application of integrated residual objects and application of nodal BC constraints? I'm surprised I haven't noticed this before; maybe because of our brief convo with respect to the Claude review. I don't like the asymmetry with the residual. We don't check for a nodal BC vector tag when deciding whether to accumulate into the vector or not, which would allow us to avoid the intermediate vector close.

I don't know if we'll ever apply Kokkos to contact, but we do rely on the intermediate population of residuals and Jacobians by integrated residual objects as then the residual at the interface represents forces which we sometimes then move over to the non-constrained surface of the contact interface. Skipping the intermediate population of the global residual/matrix would make it impossible to do that kind of kinematic constraint enforcement

This becomes very tricky in parallel. If an elemental kernel contributes to nodes covered by nodal BCs but not belonging to the current process, the remote process can't zero the nodes. I can't retrieve the closed matrix back into the COO ordering that I used for preallocation and assembly. PETSc doesn't have such capability.

That sucks. Is the COO ordering required for things like nodal bcs?

MatSetPreallocationCOO and MatSetValuesCOO are the only GPU matrix assembly method they support.

I wonder if they can make something like MatGetValuesCOO. Then I can compute nodal BCs on it.

What is different between MatZeroRows (which is what we've always used in MOOSE) and MatZeroRowsColumns with respect to the COO format?

According to the documentation:

Unlike MatZeroRows(), this ignores the MAT_KEEP_NONZERO_PATTERN option value set with MatSetOption(), it merely zeros those entries in the matrix, but never removes them from the nonzero pattern. The nonzero pattern of the matrix can still change if a nonzero needs to be inserted on a
diagonal entry that was previously missing.

I'm already zeroing out the row where nodal BCs apply in the current code I think?

So is the plan then to remove these nodal bc matrix tags, close the matrix after element assembly, apply MatZeroRowsAndColumns, and then do the nodal BC matrix application? To me I still feel like MatZeroRows with MAT_KEEP_NONZERO_PATTERN would still be a more robust pattern as in general we don't know if a nodal BC has dof coupling in which case the preset trick won't work

I implemented something. I hope the performance is not too bad because MatZeroRows is a CPU API and MatSetValuseCOO is now called twice.

lindsayad · 2026-05-13T16:51:02Z

-  Real3 grad = 0;
+  Real3 grad;


I don't like this change combined with the += below. I guess it must work but a developer who sees this would be worried that we're doing += with unitialized memory

lindsayad · 2026-05-13T16:51:17Z

+  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
+  auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
+
+  ADReal3 grad;


same concern as above

lindsayad · 2026-05-13T17:08:35Z

+    input = 'kokkos_2d_ad_diffusion_test.i'
+    cli_args = 'Outputs/exodus=false'
+    recover = false
+    requirement = 'The Kokkos system shall provide exact Jacobian using automatic differentation.'


Suggested change

requirement = 'The Kokkos system shall provide exact Jacobian using automatic differentation.'

requirement = 'The Kokkos system shall provide an exact Jacobian using automatic differentation.'

lindsayad · 2026-05-13T17:09:15Z

+    run_sim = True
+    ratio_tol = 1e-7
+    difference_tol = 1e-6
+    requirement = 'The Jacobian from KokkosADTimeDerivative shall be perfect.'


Suggested change

requirement = 'The Jacobian from KokkosADTimeDerivative shall be perfect.'

requirement = 'The Jacobian from a Kokkos based time derivative kernel shall be perfect.'

lindsayad · 2026-05-13T17:16:36Z

+    run_sim = True
+    ratio_tol = 1e-7
+    difference_tol = 1e-6
+    requirement = 'The Jacobian from KokkosADCoupledTimeDerivative shall be perfect.'


Suggested change

requirement = 'The Jacobian from KokkosADCoupledTimeDerivative shall be perfect.'

requirement = 'The Jacobian from a kokkos-based coupled variable time derivative kernel shall be perfect.'

…b#30655

NamjaeChoi · 2026-05-14T18:21:52Z

could we maybe build some unit tests to bring down the 140 unused lines introduced in Coupleable?

No, this requires the full problem to be built for testing. Hopefully we can get more coverage as the time goes on and more tests get added.

moosebuild · 2026-05-15T00:29:52Z

Job Test, step Results summary on a2f3df9 wanted to post the following:

Framework test summary

Compared against 54450aa in job civet.inl.gov/job/3821125.

Added tests

Test	Time (s)	Memory (MB)
`kokkos/kernels/ad_time_derivative.test_jacobian`	1.44	126.46
`kokkos/kernels/ad_time_derivative.coupled_jacobian`	0.92	122.19
`kokkos/kernels/ad_time_derivative.test`	0.72	114.28
`kokkos/bcs/ad_coupled_var_neumann.nonlinear/jac`	0.69	116.36
`kokkos/kernels/ad_time_derivative.coupled`	0.65	77.65
`kokkos/bcs/ad_coupled_var_neumann.test`	0.64	62.82
`kokkos/kernels/ad_diffusion.jacobian`	0.61	112.65
`kokkos/kernels/ad_time_derivative.separate`	0.60	127.85
`kokkos/bcs/ad_coupled_var_neumann.nonlinear/exo`	0.59	102.80
`kokkos/kernels/ad_diffusion.neumann`	0.57	64.75
`kokkos/kernels/ad_diffusion.dirichlet`	0.52	50.39

Modules test summary

Compared against 54450aa in job civet.inl.gov/job/3821125.

No change

NamjaeChoi changed the title ~~[WIP] Kokkos Automatic Differentation~~ [WIP] Kokkos Automatic Differentiation (Undisplaced) Apr 1, 2026

NamjaeChoi force-pushed the kokkos_ad branch from e93c29e to 39494d9 Compare April 1, 2026 23:30

NamjaeChoi added this to NEAMS MP CONNECT/GPU 26 Apr 2, 2026

NamjaeChoi moved this to In progress in NEAMS MP CONNECT/GPU 26 Apr 2, 2026

NamjaeChoi force-pushed the kokkos_ad branch 2 times, most recently from 1e1ea85 to 9e48e13 Compare April 2, 2026 20:01

NamjaeChoi changed the title ~~[WIP] Kokkos Automatic Differentiation (Undisplaced)~~ Kokkos Automatic Differentiation (Undisplaced) Apr 2, 2026

NamjaeChoi force-pushed the kokkos_ad branch 4 times, most recently from e09c807 to 50259af Compare April 3, 2026 15:30

GiudGiud assigned NamjaeChoi Apr 7, 2026

NamjaeChoi force-pushed the kokkos_ad branch from 50259af to b1cf3d9 Compare April 7, 2026 22:05

idaholab deleted a comment from moosebuild Apr 13, 2026

NamjaeChoi force-pushed the kokkos_ad branch 4 times, most recently from be28b83 to c18bb1b Compare April 21, 2026 21:53

idaholab deleted a comment from moosebuild Apr 21, 2026

NamjaeChoi force-pushed the kokkos_ad branch from c18bb1b to 64f43a0 Compare April 22, 2026 04:01

NamjaeChoi marked this pull request as ready for review April 22, 2026 04:01

NamjaeChoi requested review from joshuahansel, lindsayad and roystgnr as code owners April 22, 2026 04:01

NamjaeChoi force-pushed the kokkos_ad branch 3 times, most recently from 17b218d to a8ac7f0 Compare April 22, 2026 14:49

NamjaeChoi requested a review from cticenhour as a code owner April 22, 2026 14:49

NamjaeChoi force-pushed the kokkos_ad branch from a8ac7f0 to 98997f5 Compare April 22, 2026 16:39

moosebuild added the PR: Failed but allowed label May 11, 2026

NamjaeChoi force-pushed the kokkos_ad branch from 290161b to 58b9b09 Compare May 11, 2026 23:18

moosebuild removed the PR: Failed but allowed label May 11, 2026

NamjaeChoi force-pushed the kokkos_ad branch 2 times, most recently from b9b4ccf to e5a0d29 Compare May 11, 2026 23:37

moosebuild added the PR: Failed but allowed label May 12, 2026

lindsayad reviewed May 12, 2026

View reviewed changes

lindsayad reviewed May 13, 2026

View reviewed changes

NamjaeChoi added 10 commits May 13, 2026 15:22

Enable Kokkos MetaPhysicL idaholab#32699

8b651ba

Support computing residual and Jacobian together idaholab#32699

cc8402e

Setup for Kokkos AD idaholab#32699

32e6a1f

Add tests idaholab#32699

403f953

Add object markdowns idaholab#32699

cc868cb

Update documentations idaholab#32699

5fe0dee

Unify residual object registration macro idaholab#30655

ac83515

Remove datum.reinit() idaholab#30655

5628a53

Pad jagged array data by one to avoid reading the end of data idahola…

0a1f59a

…b#30655

Address comments 1 idaholab#32699

a5d04c7

NamjaeChoi force-pushed the kokkos_ad branch from e5a0d29 to 9204ad9 Compare May 14, 2026 18:17

moosebuild removed the PR: Failed but allowed label May 14, 2026

NamjaeChoi force-pushed the kokkos_ad branch from 9204ad9 to 0221f08 Compare May 14, 2026 18:27

Address comments 2 idaholab#32699

5e78dd3

NamjaeChoi force-pushed the kokkos_ad branch from ae7ec2f to 6997568 Compare May 14, 2026 19:17

NamjaeChoi added 2 commits May 14, 2026 17:18

Do a separate assembly of Kokkos nodal BCs idaholab#32699

8f52a02

Improve Kokkos documentations idaholab#30655

a2f3df9

NamjaeChoi force-pushed the kokkos_ad branch from 6997568 to a2f3df9 Compare May 14, 2026 23:18

	AD kernels requires `computeQpResidual()` to be defined with the following signature, where everything remains the same with the ordinary kernels except the return type being `Moose::Kokkos::ADReal`:
	AD kernels require `computeQpResidual()` to be defined with the following signature, where everything remains the same with the ordinary kernels except the return type being `Moose::Kokkos::ADReal`:

	`computeQpJacobian()` and `computeQpOffDiagJacobian()` are unused, as AD automatically assembles Jacobian.
	`computeQpJacobian()` and `computeQpOffDiagJacobian()` are unused, as AD automatically assembles the Jacobian.

	requirement = 'The Kokkos system shall provide exact Jacobian using automatic differentation.'
	requirement = 'The Kokkos system shall provide an exact Jacobian using automatic differentation.'

	requirement = 'The Jacobian from KokkosADTimeDerivative shall be perfect.'
	requirement = 'The Jacobian from a Kokkos based time derivative kernel shall be perfect.'

Conversation

NamjaeChoi commented Apr 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

moosebuild commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NamjaeChoi commented May 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lindsayad commented May 13, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NamjaeChoi May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NamjaeChoi May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NamjaeChoi May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NamjaeChoi May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NamjaeChoi commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

moosebuild commented May 15, 2026

Framework test summary

Added tests

Modules test summary

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

NamjaeChoi commented Apr 1, 2026 •

edited

Loading

moosebuild commented Apr 2, 2026 •

edited

Loading

NamjaeChoi May 13, 2026 •

edited

Loading

NamjaeChoi May 13, 2026 •

edited

Loading

NamjaeChoi May 14, 2026 •

edited

Loading

NamjaeChoi May 14, 2026 •

edited

Loading

NamjaeChoi commented May 14, 2026 •

edited

Loading