Date: 2022-03-31T04:47:20+00:00

Actual source code: ex5cu.cu

      1: static char help[] = "Test of CUDA matrix assemble with simple matrix.\n\n";

      3: // This a minimal example of the use of the CUDA MatAIJ metadata for assembly.
      4: //
      5: // The matrix must be a type 'aijcusparse' and must first be assembled on the CPU to provide the nonzero pattern.
      6: // Next, get a pointer to a simple CSR mirror (PetscSplitCSRDataStructure) of the matrix data on
      7: //    the GPU with MatCUSPARSEGetDeviceMatWrite().
      8: // Then use this object to populate the matrix on the GPU with MatSetValuesDevice().
      9: // Finally call MatAssemblyBegin/End() and the matrix is ready to use on the GPU without matrix data movement between the
     10: //    host and GPU.

     12: #include <petscconf.h>
     13: #include <petscmat.h>
     14: #include <petscdevice.h>
     15: #include <assert.h>

     17: #include <petscaijdevice.h>
     18: __global__
     19: void assemble_on_gpu(PetscSplitCSRDataStructure d_mat, PetscInt start, PetscInt end, PetscInt N, PetscMPIInt rank)
     20: {
     21:   const PetscInt  inc = blockDim.x, my0 = threadIdx.x;
     22:   PetscInt        i;
     23:   PetscErrorCode  ierr;

     25:   for (i=start+my0; i<end+1; i+=inc) {
     26:     PetscInt    js[] = {i-1, i}, nn = (i==N) ? 1 : 2; // negative indices are igored but >= N are not, so clip end
     27:     PetscScalar values[] = {1,1,1,1};
     28:     MatSetValuesDevice(d_mat,nn,js,nn,js,values,ADD_VALUES);if (ierr) assert(0);
     29:   }
     30: }

     32: PetscErrorCode assemble_on_cpu(Mat A, PetscInt start, PetscInt end, PetscInt N, PetscMPIInt rank)
     33: {
     35:   for (PetscInt i=start; i<end+1; i++) {
     36:     PetscInt    js[] = {i-1, i}, nn = (i==N) ? 1 : 2;
     37:     PetscScalar values[] = {1,1,1,1};
     38:     MatSetValues(A,nn,js,nn,js,values,ADD_VALUES);
     39:   }
     40:   return 0;
     41: }

     43: int main(int argc,char **args)
     44: {
     45:   Mat                        A;
     46:   PetscInt                   N=11, nz=3, Istart, Iend, num_threads = 128;
     47:   PetscSplitCSRDataStructure d_mat;
     48:   PetscLogEvent              event;
     49:   PetscMPIInt                rank,size;
     50:   PetscBool                  testmpiseq = PETSC_FALSE;
     51:   Vec                        x,y;

     53:   PetscInitialize(&argc,&args,(char*)0,help);
     54:   PetscOptionsGetInt(NULL,NULL, "-n", &N, NULL);
     55:   PetscOptionsGetInt(NULL,NULL, "-num_threads", &num_threads, NULL);
     56:   PetscOptionsGetInt(NULL,NULL, "-nz_row", &nz, NULL);
     57:   PetscOptionsGetBool(NULL,NULL, "-testmpiseq", &testmpiseq, NULL);
     58:   if (nz<3)   nz=3;
     59:   if (nz>N+1) nz=N+1;
     60:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
     61:   MPI_Comm_size(PETSC_COMM_WORLD,&size);

     63:   PetscLogEventRegister("GPU operator", MAT_CLASSID, &event);
     64:   MatCreateAIJCUSPARSE(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,N,N,nz,NULL,nz-1,NULL,&A);
     65:   MatSetFromOptions(A);
     66:   MatSetOption(A,MAT_IGNORE_OFF_PROC_ENTRIES,PETSC_TRUE);
     67:   MatCreateVecs(A,&x,&y);
     68:   MatGetOwnershipRange(A,&Istart,&Iend);
     69:   /* current GPU assembly code does not support offprocessor values insertion */
     70:   assemble_on_cpu(A, Istart, Iend, N, rank);
     71:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
     72:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);

     74:   // test
     75:   VecSet(x,1.0);
     76:   MatMult(A,x,y);
     77:   VecViewFromOptions(y,NULL,"-ex5_vec_view");

     79:   if (testmpiseq && size == 1) {
     80:     MatConvert(A,MATSEQAIJ,MAT_INPLACE_MATRIX,&A);
     81:     MatConvert(A,MATMPIAIJCUSPARSE,MAT_INPLACE_MATRIX,&A);
     82:   }
     83:   PetscLogEventBegin(event,0,0,0,0);
     84:   MatCUSPARSEGetDeviceMatWrite(A,&d_mat);
     85:   assemble_on_gpu<<<1,num_threads>>>(d_mat, Istart, Iend, N, rank);
     86:   cudaDeviceSynchronize();
     87:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
     88:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
     89:   PetscLogEventEnd(event,0,0,0,0);

     91:   // test
     92:   VecSet(x,1.0);
     93:   MatMult(A,x,y);
     94:   VecViewFromOptions(y,NULL,"-ex5_vec_view");

     96:   MatDestroy(&A);
     97:   VecDestroy(&x);
     98:   VecDestroy(&y);
     99:   PetscFinalize();
    100:   return 0;
    101: }

    103: /*TEST

    105:    build:
    106:       requires: cuda

    108:    test:
    109:       suffix: 0
    110:       diff_args: -j
    111:       args: -n 11 -ex5_vec_view
    112:       nsize: 1

    114:    test:
    115:       suffix: 1
    116:       diff_args: -j
    117:       args: -n 11 -ex5_vec_view
    118:       nsize: 2

    120:    test:
    121:       suffix: 2
    122:       diff_args: -j
    123:       args: -n 11 -testmpiseq -ex5_vec_view
    124:       nsize: 1

    126: TEST*/