OffloadPC (CUDA GPU) #3784
OffloadPC (CUDA GPU) #3784
Conversation
|
|
There was a problem hiding this comment.
Delete this file as discussed.
| @@ -55,6 +55,8 @@ def __init__(self, A, *, P=None, solver_parameters=None, | |||
| solver_parameters = solving_utils.set_defaults(solver_parameters, | |||
| A.arguments(), | |||
| ksp_defaults=self.DEFAULT_KSP_PARAMETERS) | |||
| # todo: add offload to solver parameters - how? prefix? | |||
There was a problem hiding this comment.
This needs addressing somehow. Not quite sure what is meant by this comment.
| # u.getArray() | ||
|
|
||
| # else: | ||
| # instead: preconditioner |
| from firedrake.petsc import PETSc | ||
| from firedrake.ufl_expr import TestFunction, TrialFunction | ||
| import firedrake.dmhooks as dmhooks | ||
| from firedrake.dmhooks import get_function_space |
There was a problem hiding this comment.
Probably an unnecessary import.
|
|
||
|
|
||
| class OffloadPC(PCBase): | ||
| """Offload PC from CPU to GPU and back. |
There was a problem hiding this comment.
This docstring could perhaps contain more detail about what is actually happening and even perhaps why one may wish to do this.
There was a problem hiding this comment.
E.g. This is only for CUDA GPUs
| x_cu = PETSc.Vec() | ||
| x_cu.createCUDAWithArrays(x) # end | ||
| with PETSc.Log.Event("Event: solve"): | ||
| self.pc.apply(x_cu, y_cu) # |
There was a problem hiding this comment.
| self.pc.apply(x_cu, y_cu) # | |
| self.pc.apply(x_cu, y_cu) |
| with PETSc.Log.Event("Event: solve"): | ||
| self.pc.apply(x_cu, y_cu) # | ||
| with PETSc.Log.Event("Event: vectors copy back"): | ||
| y.copy(y_cu) # |
There was a problem hiding this comment.
| y.copy(y_cu) # | |
| y.copy(y_cu) |
| y.copy(y_cu) # | ||
|
|
||
| def applyTranspose(self, pc, X, Y): | ||
| raise NotImplementedError |
There was a problem hiding this comment.
Maybe have a useful error message? Not sure what the usual approach is here.
|
|
||
| def view(self, pc, viewer=None): | ||
| super().view(pc, viewer) | ||
| print("viewing PC") |
There was a problem hiding this comment.
| print("viewing PC") |
| if not isinstance(A, firedrake.matrix.AssembledMatrix): | ||
| # linear MG doesn't need RHS, supply zero. | ||
| lvp = vs.LinearVariationalProblem(a=A.a, L=0, u=x, bcs=A.bcs) | ||
| mat_type = A.mat_type |
There was a problem hiding this comment.
pointless line to have, delete and use A.mat_type below
Description
This adds an offloading preconditioner.