Offloading CUDA #4166
Offloading CUDA #4166
Conversation
firedrake/preconditioners/offload.py
Outdated
| __all__ = ("OffloadPC",) | ||
|
|
||
|
|
||
| class OffloadPC(PCBase): |
There was a problem hiding this comment.
Could AssembledPC assume this functionality by providing -assembled_mat_type aijcusparse?
There was a problem hiding this comment.
Possibly, but then you lose all flexibility w.r.t. using other matrix types.
There was a problem hiding this comment.
We could have this one as a subclass of AssembledPC, there's substantial code duplication
There was a problem hiding this comment.
That seems like a good idea. @Olender
There was a problem hiding this comment.
Good idea! I implemented this in the latest commit, but I'm still testing things out to ensure everything works as expected. Let me know if you have any more suggestions
|
By the way this shouldn't be named "DO NOT MERGE". Just leave it as a draft PR. |
.github/workflows/build_cuda.yml
Outdated
| . venv/bin/activate | ||
| : # Use pytest-xdist here so we can have a single collated output (not possible | ||
| : # for parallel tests) | ||
| firedrake-run-split-tests 1 1 "-n 12 $EXTRA_PYTEST_ARGS" |
There was a problem hiding this comment.
I think this is failing because EXTRA_PYTEST_ARGS is not defined.
firedrake/preconditioners/offload.py
Outdated
| # We set a DM and an appropriate SNESContext on the constructed PC | ||
| # so one can do e.g. multigrid or patch solves. | ||
| dm = outer_pc.getDM() | ||
| self._ctx_ref = self.new_snes_ctx( | ||
| outer_pc, a, bcs, mat_type, | ||
| fcp=fcp, options_prefix=options_prefix | ||
| ) |
There was a problem hiding this comment.
You might not need to create a new _SNESContext here, but instead just grab it from the parent PC. I guess most of the symbolic stuff involving the form above wouldn't be required anymore, so this will look very different from AssembledPC
| # Update preconditioner with GPU matrix | ||
| self.pc.setOperators(A, P_cu) | ||
|
|
||
| def form(self, pc, test, trial): |
There was a problem hiding this comment.
This method is already inherited from PCBase
Description