Skip to content

Offloading CUDA#4166

Closed
Olender wants to merge 49 commits into
mainfrom
olender/CUDA
Closed

Offloading CUDA#4166
Olender wants to merge 49 commits into
mainfrom
olender/CUDA

Conversation

@Olender
Copy link
Copy Markdown
Contributor

@Olender Olender commented Mar 27, 2025

Description

Comment thread firedrake/preconditioners/offload.py Outdated
__all__ = ("OffloadPC",)


class OffloadPC(PCBase):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could AssembledPC assume this functionality by providing -assembled_mat_type aijcusparse?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly, but then you lose all flexibility w.r.t. using other matrix types.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could have this one as a subclass of AssembledPC, there's substantial code duplication

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That seems like a good idea. @Olender

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@connorjward
Copy link
Copy Markdown
Contributor

By the way this shouldn't be named "DO NOT MERGE". Just leave it as a draft PR.

@Olender Olender changed the title DO NOT MERGE: Offloading CUDA Offloading CUDA Apr 4, 2025
Comment thread .github/workflows/build_cuda.yml Outdated
Comment thread .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"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is failing because EXTRA_PYTEST_ARGS is not defined.

Comment thread firedrake/preconditioners/offload.py Outdated
Comment on lines +65 to +71
# 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
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method is already inherited from PCBase

@connorjward
Copy link
Copy Markdown
Contributor

This was merged in #4953

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants