mpi/kokkos/mpiaijkok.kokkos.cxx

124   MatScalarKokkosView rootBuf, leafBuf;
149   PetscIntKokkosView irootloc; // plans to put E (i.e., Bd, Bo) into rootBuf
154   MatColIdxKokkosView Fdjmap; // plans to reduce data in rootBuf to Fd, Fo
600     mm->rootBuf  = MatScalarKokkosView(NoInit("rootBuf"), nroots);
637   auto       &rootBuf      = mm->rootBuf;
645   // Copy rows in A/B of E to leafBuf, then pass it to rootBuf
669   PetscCall(PetscSFReduceWithMemTypeBegin(reduceSF, MPIU_SCALAR, PETSC_MEMTYPE_KOKKOS, leafBuf.data(), PETSC_MEMTYPE_KOKKOS, rootBuf.data(), MPI_REPLACE));
677   auto       &rootBuf  = mm->rootBuf;
689   PetscCall(PetscSFReduceEnd(reduceSF, MPIU_SCALAR, leafBuf.data(), rootBuf.data(), MPI_REPLACE));
691   // Reduce data in rootBuf to Fd and Fo
695       for (MatRowMapType k = Fdjmap(i); k < Fdjmap(i + 1); k++) sum += rootBuf(Fdjperm(k));
702       for (MatRowMapType k = Fojmap(i); k < Fojmap(i + 1); k++) sum += rootBuf(Fojperm(k));
918     mm->rootBuf   = MatScalarKokkosView(NoInit("rootBuf"), nroots);
937     PetscCall(MatMergeGetLaunchParameters<DefaultExecutionSpace>(mm->irootloc.extent(0), mm->rootBuf.extent(0), -1, teamSize, vectorLength, rowsPerTeam));
962   MatScalarKokkosView &rootBuf   = mm->rootBuf;
972   // Copy rows in A/B of E to rootBuf, then bcast it to leafBuf
986               rootBuf(disp + j) = Ba(Bi(i) + j);
988               rootBuf(disp + j) = Aa(Ai(i) + j - nzleft);
990               rootBuf(disp + j) = Ba(Bi(i) + j - alen);
996   PetscCall(PetscSFBcastWithMemTypeBegin(bcastSF, MPIU_SCALAR, PETSC_MEMTYPE_KOKKOS, rootBuf.data(), PETSC_MEMTYPE_KOKKOS, leafBuf.data(), MPI_REPLACE));
1014   MatScalarKokkosView &rootBuf      = mm->rootBuf;
1021   PetscCall(PetscSFBcastEnd(bcastSF, MPIU_SCALAR, rootBuf.data(), leafBuf.data(), MPI_REPLACE));