Skip to content

Minor updates to OffloadPC#5164

Open
dsroberts wants to merge 3 commits into
firedrakeproject:mainfrom
dsroberts:dsroberts/minor-offload-fixes
Open

Minor updates to OffloadPC#5164
dsroberts wants to merge 3 commits into
firedrakeproject:mainfrom
dsroberts:dsroberts/minor-offload-fixes

Conversation

@dsroberts

Copy link
Copy Markdown
Contributor

Description

Two minor updates to OffloadPC based on continued local testing. The first change is preventing certain types of of PETSc matrix from being offloaded. When implicit matrices are converted to aijcusparse matrices, PETSc constructs a dense matrix row-by-row then converts that to a sparse matrix. Since PETSc can transparently manage host-device transfers when using non-aijcusparse matrix types on cuda vectors, it is far more efficient for PETSc to offload as needed than to do a full implicit to dense to sparse conversion. I'm not sure this is a complete list of implicit matrix types, but these are the types I run into with our testing.

The second change is in the construction of the device matrices. If the out kwarg is not specified or set to None in Mat.convert(), the conversion happens in place meaning e.g. both P and P_device are now the same aijcusparse matrix. This was probably not picked up earlier as an aijcusparse matrix is an aij matrix with a managed device buffer under the hood, so silently converting an aij matrix probably does not affect any other part of a job, but it might cause issues with other matrix types.

There is a check to make sure the silent conversion hasn't happened, but I don't know how to test for the implicit -> dense -> sparse conversion. The right thing to do I think would be to check that A is A_dev if A.getType() is one of the 'do not offload' types, but I don't know how to retrieve that information out of nested PCs.

@dsroberts dsroberts added the ci:cuda Run the test suite with a CUDA-enabled build label Jun 11, 2026
Comment thread firedrake/preconditioners/offload.py Outdated
Co-authored-by: Josh Hope-Collins <jhc.jss@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:cuda Run the test suite with a CUDA-enabled build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants