<div dir="ltr"><div>Perfect, that seems to have fixed the issue.</div><div><br></div><div>Thanks for your help!</div><div>Steven</div><br></div><br><div class="gmail_quote gmail_quote_container"><div dir="ltr" class="gmail_attr">On Thu, 3 Apr 2025 at 04:52, Junchao Zhang <<a href="mailto:junchao.zhang@gmail.com">junchao.zhang@gmail.com</a>> wrote:<br></div><blockquote class="gmail_quote" style="margin:0px 0px 0px 0.8ex;border-left:1px solid rgb(204,204,204);padding-left:1ex"><div dir="ltr"><div>Hi, Steven,</div><div> Thanks for the test, which helped me easily find the petsc bug. I have a fix at <a href="https://urldefense.us/v3/__https://gitlab.com/petsc/petsc/-/merge_requests/8272__;!!G_uCfscf7eWS!ZNkvUNeiuBo6kRAR5k_yCtp_mOCmBVppGB_yvYsMwSiohil2UIhWXsli1sEeYZvqM0PkQe6HjVYqR8YxD-S1fT_OHcxzNRK2$" target="_blank">https://gitlab.com/petsc/petsc/-/merge_requests/8272</a>.</div><div> VecKokkos does not use offloadmask for its gpu/cpu sync state, while VecCUDA/VecHIP do. My expectation is users should not use VecGetOffloadMask(), because it is too low level. We have bad API design here.</div><div> </div><div> Thank you!</div><div><div dir="ltr" class="gmail_signature"><div dir="ltr">--Junchao Zhang</div></div></div><br></div><br><div class="gmail_quote"><div dir="ltr" class="gmail_attr">On Wed, Apr 2, 2025 at 7:11 PM Steven Dargaville <<a href="mailto:dargaville.steven@gmail.com" target="_blank">dargaville.steven@gmail.com</a>> wrote:<br></div><blockquote class="gmail_quote" style="margin:0px 0px 0px 0.8ex;border-left:1px solid rgb(204,204,204);padding-left:1ex"><div dir="ltr"><div>Hi</div><div><br></div><div>I have some code that does a solve with a PCMAT preconditioner. The mat used is a shell and inside the shell MatMult it calls VecPointwiseDivide with a vector "diag" that is the diagonal of a matrix assigned outside the shell. <br></div><div><br></div><div>If I use mat/vec type of cuda, this occurs without any gpu/cpu copies as I would expect. If however I use mat/vec type kokkos, at every iteration of the solve there is a gpu/cpu copy that occurs. It seems this is triggered by the offloadmask in the vector "diag", as it stays as 1 and hence a copy occurs in VecPointwiseDivide.</div><div><br></div><div>I would have expected the offload mask to be 256 (kokkos) after the first iteration, as the offload mask of "diag" changes to 3 when using cuda after the first iteration. <br></div><div><br></div><div>Is this the expected behaviour with Kokkos, or is there something I need to do to trigger that "diag" has its values on the gpu to prevent copies? I have example c++ code that demonstrates this below. You can see the difference when run with petsc 3.23.0 and either "-log_view -mat_type aijcusparse -vec_type cuda" or "-log_view -mat_type aijkokkos -vec_type kokkos".</div><div><br></div><div>Thanks for your help</div><div>Steven</div><div><br></div><div>Example c++ code:</div><div><br></div><div>// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~</div><div><br></div>static char help[] = "Tests Kokkos for SHELL matrices\n\n";<br><br>#include <iostream><br>#include <petscksp.h><br>#include <petsclog.h><br><br>typedef struct _n_User *User;<br>struct _n_User {<br> Mat A;<br> Vec diag;<br>};<br><br>static PetscErrorCode MatMult_User(Mat A, Vec X, Vec Y)<br>{<br> User user;<br><br> PetscFunctionBegin;<br> PetscCall(MatShellGetContext(A, &user));<br><br> // Print the offload mask inside the matmult<br> PetscOffloadMask offloadmask;<br> PetscCall(VecGetOffloadMask(X, &offloadmask));<br> std::cout << "offload inside X " << offloadmask << std::endl;<br> PetscCall(VecGetOffloadMask(Y, &offloadmask));<br> std::cout << "offload inside Y " << offloadmask << std::endl;<br> PetscCall(VecGetOffloadMask(user->diag, &offloadmask));<br> std::cout << "offload inside diag " << offloadmask << std::endl; <br><br> PetscCall(VecPointwiseDivide(Y, X, user->diag));<br> PetscFunctionReturn(PETSC_SUCCESS);<br>}<br><br>int main(int argc, char **args)<br>{<br> const PetscScalar xvals[] = {11, 13}, yvals[] = {17, 19};<br> const PetscInt inds[] = {0, 1};<br> PetscScalar avals[] = {2, 3, 5, 7};<br> Mat S1, A;<br> Vec X, Y, diag;<br> KSP ksp; <br> PC pc;<br> User user;<br> PetscLogStage stage1, gpu_copy;<br> <br> PetscFunctionBeginUser;<br> PetscCall(PetscInitialize(&argc, &args, NULL, help));<br><br> // Build a matrix and vectors<br> PetscCall(MatCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, 2, 2, 2, 2, &A));<br> PetscCall(MatSetUp(A));<br> PetscCall(MatSetValues(A, 2, inds, 2, inds, avals, INSERT_VALUES));<br> PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));<br> PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));<br> PetscCall(MatCreateVecs(A, NULL, &X));<br> PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, 2, 2, &X));<br> PetscCall(VecSetValues(X, 2, inds, xvals, INSERT_VALUES));<br> PetscCall(VecDuplicate(X, &Y));<br> PetscCall(VecDuplicate(X, &diag));<br> PetscCall(VecSetValues(Y, 2, inds, yvals, INSERT_VALUES));<br> PetscCall(VecAssemblyBegin(Y));<br> PetscCall(VecAssemblyEnd(Y));<br><br> // Create a shell matrix<br> PetscCall(MatGetDiagonal(A, diag));<br> PetscCall(PetscNew(&user));<br> user->A = A;<br> user->diag = diag;<br> PetscCall(MatCreateShell(PETSC_COMM_WORLD, 2, 2, 2, 2, user, &S1));<br> PetscCall(MatSetUp(S1));<br> PetscCall(MatShellSetOperation(S1, MATOP_MULT, (void (*)(void))MatMult_User));<br> PetscCall(MatAssemblyBegin(S1, MAT_FINAL_ASSEMBLY));<br> PetscCall(MatAssemblyEnd(S1, MAT_FINAL_ASSEMBLY));<br> <br> // Do a solve<br> PetscCall(KSPCreate(PETSC_COMM_WORLD,&ksp));<br> // Give the ksp a pcmat as the preconditioner and the mat is the shell<br> PetscCall(KSPSetOperators(ksp,A, S1));<br> PetscCall(KSPSetType(ksp, KSPRICHARDSON));<br> PetscCall(KSPSetFromOptions(ksp));<br> PetscCall(KSPGetPC(ksp, &pc));<br> PetscCall(PCSetType(pc, PCMAT));<br> PetscCall(KSPSetUp(ksp));<br><br> // Print the offload mask before our solve<br> PetscOffloadMask offloadmask;<br> PetscCall(VecGetOffloadMask(X, &offloadmask));<br> std::cout << "offload X " << offloadmask << std::endl;<br> PetscCall(VecGetOffloadMask(Y, &offloadmask));<br> std::cout << "offload Y " << offloadmask << std::endl;<br> PetscCall(VecGetOffloadMask(user->diag, &offloadmask));<br> std::cout << "offload diag " << offloadmask << std::endl; <br><br> // Trigger any gpu copies in the first solve<br> PetscCall(PetscLogStageRegister("gpu_copy",&gpu_copy));<br> PetscCall(PetscLogStagePush(gpu_copy));<br> PetscCall(KSPSolve(ksp, X, Y));<br> PetscCall(PetscLogStagePop()); <br><br> // There should be no copies in this solve<br> PetscCall(PetscLogStageRegister("no copy",&stage1));<br> PetscCall(PetscLogStagePush(stage1)); <br> PetscCall(KSPSolve(ksp, X, Y));<br> PetscCall(PetscLogStagePop());<br><br> PetscCall(MatDestroy(&S1));<br> PetscCall(VecDestroy(&X));<br> PetscCall(VecDestroy(&Y));<br> PetscCall(PetscFinalize());<br> return 0;<br>}</div>
</blockquote></div>
</blockquote></div>