[petsc-users] building kokkos matrices on the device

[Steven Dargaville]

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.zhang at 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.steven at 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
>>
>
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.mcs.anl.gov/pipermail/petsc-users/attachments/20250225/88ea1a14/attachment-0001.html>