1: static char help[] = "Test MatDenseGetSubMatrix() on a CUDA matrix and MatSetInf() on dense matrices.\n";

  3: #include <petscmat.h>

  5: /*
  6:    Every entry of the local block must be positive infinity, while the padding between two columns, which is only present when the leading dimension is larger than the number of local rows, must be left alone
  7: */
  8: static PetscErrorCode CheckInf(Mat A)
  9: {
 10:   const PetscScalar *a;
 11:   PetscInt           i, j, m, N, lda;

 13:   PetscFunctionBeginUser;
 14:   PetscCall(MatGetLocalSize(A, &m, NULL));
 15:   PetscCall(MatGetSize(A, NULL, &N));
 16:   PetscCall(MatDenseGetLDA(A, &lda));
 17:   PetscCall(MatDenseGetArrayRead(A, &a));
 18:   for (j = 0; j < N; j++) {
 19:     for (i = 0; i < m; i++) PetscCheck(PetscIsInfReal(PetscRealPart(a[i + j * lda])) && PetscRealPart(a[i + j * lda]) > 0.0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Entry (%" PetscInt_FMT ",%" PetscInt_FMT ") is not positive infinity", i, j);
 20:     for (i = m; i < lda; i++) PetscCheck(a[i + j * lda] == 0.0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Padding entry (%" PetscInt_FMT ",%" PetscInt_FMT ") was overwritten", i, j);
 21:   }
 22:   PetscCall(MatDenseRestoreArrayRead(A, &a));
 23:   PetscFunctionReturn(PETSC_SUCCESS);
 24: }

 26: /*
 27:    Device dense Mat with n local rows, n global columns and the given leading dimension. MatCreateDenseCUDA() preallocates the storage before MatDenseSetLDA() can take effect, hence the MatSetType() path below
 28: */
 29: static PetscErrorCode CreateDenseCUDA(PetscInt n, PetscInt lda, Mat *A)
 30: {
 31:   PetscFunctionBeginUser;
 32:   PetscCall(MatCreate(PETSC_COMM_WORLD, A));
 33:   PetscCall(MatSetSizes(*A, n, PETSC_DECIDE, PETSC_DETERMINE, n));
 34:   PetscCall(MatSetType(*A, MATDENSECUDA));
 35:   PetscCall(MatDenseSetLDA(*A, lda));
 36:   PetscCall(MatSetUp(*A));
 37:   PetscFunctionReturn(PETSC_SUCCESS);
 38: }

 40: /*
 41:    MatSetInf() on a dense Mat whose leading dimension matches the number of local rows and on one whose leading dimension is larger, on the host or on the device
 42: */
 43: static PetscErrorCode TestSetInf(PetscInt n, PetscInt lda, PetscBool cuda)
 44: {
 45:   Mat          A;
 46:   PetscScalar *data = NULL;

 48:   PetscFunctionBeginUser;
 49:   if (cuda) PetscCall(CreateDenseCUDA(n, n, &A));
 50:   else PetscCall(MatCreateDense(PETSC_COMM_WORLD, n, PETSC_DECIDE, PETSC_DETERMINE, n, NULL, &A));
 51:   PetscCall(MatZeroEntries(A));
 52:   PetscCall(MatSetInf(A));
 53:   PetscCall(CheckInf(A));
 54:   PetscCall(MatDestroy(&A));

 56:   if (cuda) {
 57:     PetscCall(CreateDenseCUDA(n, lda, &A));
 58:     /* the values are put on the device so that MatSetInf() has to invalidate them, as when it is called on the block of solutions of a KSPMatSolve() which has not converged */
 59:     PetscCall(MatZeroEntries(A));
 60:   } else {
 61:     PetscCall(PetscCalloc1((size_t)lda * n, &data));
 62:     PetscCall(MatCreateDense(PETSC_COMM_WORLD, n, PETSC_DECIDE, PETSC_DETERMINE, n, data, &A));
 63:     PetscCall(MatDenseSetLDA(A, lda));
 64:   }
 65:   PetscCall(MatSetInf(A));
 66:   PetscCall(CheckInf(A));
 67:   PetscCall(MatDestroy(&A));
 68:   PetscCall(PetscFree(data));
 69:   PetscFunctionReturn(PETSC_SUCCESS);
 70: }

 72: int main(int argc, char **argv)
 73: {
 74:   Mat          A, B;
 75:   PetscScalar *b;
 76:   PetscInt     n = 4, lda = 5, i, k;
 77:   PetscBool    cuda = PETSC_FALSE, set_inf = PETSC_FALSE;

 79:   PetscFunctionBeginUser;
 80:   PetscCall(PetscInitialize(&argc, &argv, 0, help));
 81:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
 82:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-lda", &lda, NULL));
 83:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-cuda", &cuda, NULL));
 84:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-set_inf", &set_inf, NULL));
 85:   PetscCheck(lda >= n, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "lda %" PetscInt_FMT " < n %" PetscInt_FMT, lda, n);

 87:   if (set_inf) PetscCall(TestSetInf(n, lda, cuda));
 88:   else {
 89: #if PetscDefined(HAVE_CUDA)
 90:     if (cuda) PetscCall(MatCreateSeqDenseCUDA(PETSC_COMM_SELF, lda, n, NULL, &A));
 91:     else
 92: #endif
 93:       PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, lda, n, NULL, &A));

 95:     for (k = 0; k < 3; k++) {
 96:       PetscCall(MatDenseGetSubMatrix(A, 0, n, 0, n, &B));
 97:       PetscCall(MatDenseGetArray(B, &b));
 98:       for (i = 0; i < n; i++) {
 99:         b[i + i * lda] = 2.0 * (i + 1);
100:         if (i > 0) b[i + (i - 1) * lda] = (PetscReal)(k + 1);
101:       }
102:       PetscCall(MatDenseRestoreArray(B, &b));
103:       PetscCall(MatDenseRestoreSubMatrix(A, &B));
104:       PetscCall(MatView(A, NULL));
105:     }

107:     PetscCall(MatDestroy(&A));
108:   }
109:   PetscCall(PetscFinalize());
110:   return 0;
111: }

113: /*TEST

115:    testset:
116:       output_file: output/ex257_1.out
117:       diff_args: -j
118:       test:
119:          suffix: 1
120:       test:
121:          suffix: 1_cuda
122:          args: -cuda
123:          requires: cuda
124:          filter: sed -e "s/seqdensecuda/seqdense/"

126:    test:
127:       suffix: set_inf
128:       nsize: {{1 2}}
129:       args: -set_inf
130:       output_file: output/empty.out

132:    test:
133:       suffix: set_inf_cuda
134:       nsize: {{1 2}}
135:       args: -set_inf -cuda
136:       requires: cuda
137:       output_file: output/empty.out

139: TEST*/