Those two constructors would definitely meet my needs, thanks! Also I should note that the comment about garray and B in MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices is correct if garray is passed in as NULL, it's just that if you pass in a completed garray it doesn't bother creating one or changing the column indices of B. So I would suggest the comment be: "if garray is NULL the offdiag matrix B should have global col ids; if garray is not NULL the offdiag matrix B should have local col ids"
On Wed, 26 Feb 2025 at 03:35, Junchao Zhang <junchao.zh...@gmail.com> wrote: > Mat_SeqAIJKokkos is private because it is in a private header > src/mat/impls/aij/seq/kokkos/aijkok.hpp > > Your observation about the garray in > MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices() > might be right. The comment > > - B - the offdiag matrix using global col ids > > is out of date. Perhaps it should be "the offdiag matrix uses local column > indices and garray contains the local to global mapping". But I need to > double check it. > > Since you use Kokkos, I think we could provide these two constructors for > MATSEQAIJKOKKOS and MATMPIAIJKOKKOS respectively > > - MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm comm, > KokkosCsrMatrix csr, Mat *A) > > > - MatCreateMPIAIJKokkosWithSeqAIJKokkos(MPI_Comm comm, Mat A, Mat B, > PetscInt *garray, Mat *mat) > > // To mimic the existing MatCreateMPIAIJWithSeqAIJ(MPI_Comm comm, > Mat A, Mat B, const PetscInt garray[], Mat *mat); > // A and B are MATSEQAIJKOKKOS matrices and use local column > indices > > Do they meet your needs? > > --Junchao Zhang > > > On Tue, Feb 25, 2025 at 5:35 PM Steven Dargaville < > dargaville.ste...@gmail.com> wrote: > >> Thanks for the response! >> >> Although MatSetValuesCOO happens on the device if the input coo_v pointer >> is device memory, I believe MatSetPreallocationCOO requires host pointers >> for coo_i and coo_j, and the preallocation (and construction of the COO >> structures) happens on the host and is then copied onto the device? I need >> to be able to create a matrix object with minimal work on the host (like >> many of the routines in aijkok.kokkos.cxx do internally). I originally used >> the COO interface to build the matrices I need, but that was around 5x >> slower than constructing the aij structures myself on the device and then >> just directly using the MatSetSeqAIJKokkosWithCSRMatrix type methods. >> >> The reason I thought MatSetSeqAIJKokkosWithCSRMatrix could easily be made >> public is that the Mat_SeqAIJKokkos constructors are already publicly >> accessible? In particular one of those constructors takes in pointers to >> the Kokkos dual views which store a,i,j, and hence one can build a >> sequential matrix with nothing (or very little) occuring on the host. The >> only change I can see that would be necessary is for >> MatSetSeqAIJKokkosWithCSRMatrix (or MatCreateSeqAIJKokkosWithCSRMatrix) to >> be public is to change the PETSC_INTERN to PETSC_EXTERN? >> >> For MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices, I believe all that >> is required is declaring the method in the .hpp, as it's already defined as >> static in mpiaijkok.kokkos.cxx. In particular, the comments >> above MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices suggest that the >> off-diagonal block B needs to be built with global column ids, with >> mpiaij->garray constructed on the host along with the rewriting of the >> global column indices in B. This happens in MatSetUpMultiply_MPIAIJ, but >> checking the code there shows that if you pass in a non-null garray to >> MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices, the construction and >> compatification is skipped, meaning B can be built with local column ids as >> long as garray is provided on the host (which I also build on the device >> and then just copy to the host). Again this is what some of the internal >> Kokkos routines rely on, like the matrix-product. >> >> I am happy to try doing this and submitting a request to the petsc gitlab >> if this seems sensible, I just wanted to double check that I wasn't missing >> something important? >> Thanks >> Steven >> >> On Tue, 25 Feb 2025 at 22:16, Junchao Zhang <junchao.zh...@gmail.com> >> wrote: >> >>> Hi, Steven, >>> MatSetSeqAIJKokkosWithCSRMatrix(Mat A, Mat_SeqAIJKokkos *akok) uses a >>> private data type Mat_SeqAIJKokkos, so it can not be directly made public. >>> If you already use COO, then why not directly make the matrix of type >>> MATAIJKOKKOS and call MatSetPreallocationCOO() and MatSetValuesCOO()? >>> So I am confused by your needs. >>> >>> Thanks! >>> --Junchao Zhang >>> >>> >>> On Tue, Feb 25, 2025 at 3:39 PM Steven Dargaville < >>> dargaville.ste...@gmail.com> wrote: >>> >>>> Hi >>>> >>>> I'm just wondering if there is any possibility of making: >>>> MatSetSeqAIJKokkosWithCSRMatrix or MatCreateSeqAIJKokkosWithCSRMatrix >>>> in src/mat/impls/aij/seq/kokkos/aijkok.kokkos.cxx >>>> MatSetMPIAIJKokkosWithSplitSeqAIJKokkosMatrices in >>>> src/mat/impls/aij/mpi/kokkos/mpiaijkok.kokkos.cxx >>>> >>>> publicly accessible outside of petsc, or if there is an interface I >>>> have missed for creating Kokkos matrices entirely on the device? >>>> MatCreateSeqAIJKokkosWithCSRMatrix for example is marked as PETSC_INTERN so >>>> I can't link to it. >>>> >>>> I've currently just copied the code inside of those methods so that I >>>> can build without any preallocation on the host (e.g., through the COO >>>> interface) and it works really well. >>>> >>>> Thanks for your help >>>> Steven >>>> >>>
