1: #include "../denseqn.h"
  2: #include <petsc/private/cupminterface.hpp>
  3: #include <petsc/private/cupmobject.hpp>

  5: namespace Petsc
  6: {

  8: namespace device
  9: {

 11: namespace cupm
 12: {

 14: namespace impl
 15: {

 17: template <DeviceType T>
 18: struct UpperTriangular : CUPMObject<T> {
 19:   PETSC_CUPMOBJECT_HEADER(T);

 21:   static PetscErrorCode SolveInPlace(PetscDeviceContext, PetscBool, PetscInt, const PetscScalar[], PetscInt, PetscScalar[], PetscInt) noexcept;
 22:   static PetscErrorCode SolveInPlaceCyclic(PetscDeviceContext, PetscBool, PetscInt, PetscInt, PetscInt, const PetscScalar[], PetscInt, PetscScalar[], PetscInt) noexcept;
 23: };

 25: template <DeviceType T>
 26: PetscErrorCode UpperTriangular<T>::SolveInPlace(PetscDeviceContext dctx, PetscBool hermitian_transpose, PetscInt N, const PetscScalar A[], PetscInt lda, PetscScalar x[], PetscInt stride) noexcept
 27: {
 28:   cupmBlasInt_t    n;
 29:   cupmBlasHandle_t handle;
 30:   auto             A_ = cupmScalarPtrCast(A);
 31:   auto             x_ = cupmScalarPtrCast(x);

 33:   PetscFunctionBegin;
 34:   if (!N) PetscFunctionReturn(PETSC_SUCCESS);
 35:   PetscCall(PetscCUPMBlasIntCast(N, &n));
 36:   PetscCall(GetHandlesFrom_(dctx, &handle));
 37:   PetscCall(PetscLogGpuTimeBegin());
 38:   PetscCallCUPMBLAS(cupmBlasXtrsv(handle, CUPMBLAS_FILL_MODE_UPPER, hermitian_transpose ? CUPMBLAS_OP_C : CUPMBLAS_OP_N, CUPMBLAS_DIAG_NON_UNIT, n, A_, lda, x_, stride));
 39:   PetscCall(PetscLogGpuTimeEnd());

 41:   PetscCall(PetscLogGpuFlops(1.0 * N * N));
 42:   PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
 43:   PetscFunctionReturn(PETSC_SUCCESS);
 44: }

 46: template <DeviceType T>
 47: PetscErrorCode UpperTriangular<T>::SolveInPlaceCyclic(PetscDeviceContext dctx, PetscBool hermitian_transpose, PetscInt m, PetscInt oldest, PetscInt next, const PetscScalar A[], PetscInt lda, PetscScalar x[], PetscInt stride) noexcept
 48: {
 49:   PetscInt              N            = next - oldest;
 50:   PetscInt              oldest_index = oldest % m;
 51:   PetscInt              next_index   = next % m;
 52:   cupmBlasInt_t         n_old, n_new;
 53:   cupmBlasPointerMode_t pointer_mode;
 54:   cupmBlasHandle_t      handle;
 55:   auto                  sone      = cupmScalarCast(1.0);
 56:   auto                  minus_one = cupmScalarCast(-1.0);
 57:   auto                  A_        = cupmScalarPtrCast(A);
 58:   auto                  x_        = cupmScalarPtrCast(x);

 60:   PetscFunctionBegin;
 61:   if (!N) PetscFunctionReturn(PETSC_SUCCESS);
 62:   PetscCall(PetscCUPMBlasIntCast(m - oldest_index, &n_old));
 63:   PetscCall(PetscCUPMBlasIntCast(next_index, &n_new));
 64:   PetscCall(GetHandlesFrom_(dctx, &handle));
 65:   PetscCall(PetscLogGpuTimeBegin());
 66:   PetscCallCUPMBLAS(cupmBlasGetPointerMode(handle, &pointer_mode));
 67:   PetscCallCUPMBLAS(cupmBlasSetPointerMode(handle, CUPMBLAS_POINTER_MODE_HOST));
 68:   if (!hermitian_transpose) {
 69:     if (n_new > 0) PetscCallCUPMBLAS(cupmBlasXtrsv(handle, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_N, CUPMBLAS_DIAG_NON_UNIT, n_new, A_, lda, x_, stride));
 70:     if (n_new > 0 && n_old > 0) PetscCallCUPMBLAS(cupmBlasXgemv(handle, CUPMBLAS_OP_N, n_old, n_new, &minus_one, &A_[oldest_index], lda, x_, stride, &sone, &x_[oldest_index], stride));
 71:     if (n_old > 0) PetscCallCUPMBLAS(cupmBlasXtrsv(handle, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_N, CUPMBLAS_DIAG_NON_UNIT, n_old, &A_[oldest_index * (lda + 1)], lda, &x_[oldest_index], stride));
 72:   } else {
 73:     if (n_old > 0) PetscCallCUPMBLAS(cupmBlasXtrsv(handle, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_C, CUPMBLAS_DIAG_NON_UNIT, n_old, &A_[oldest_index * (lda + 1)], lda, &x_[oldest_index], stride));
 74:     if (n_new > 0 && n_old > 0) PetscCallCUPMBLAS(cupmBlasXgemv(handle, CUPMBLAS_OP_C, n_old, n_new, &minus_one, &A_[oldest_index], lda, &x_[oldest_index], stride, &sone, x_, stride));
 75:     if (n_new > 0) PetscCallCUPMBLAS(cupmBlasXtrsv(handle, CUPMBLAS_FILL_MODE_UPPER, CUPMBLAS_OP_C, CUPMBLAS_DIAG_NON_UNIT, n_new, A_, lda, x_, stride));
 76:   }
 77:   PetscCallCUPMBLAS(cupmBlasSetPointerMode(handle, pointer_mode));
 78:   PetscCall(PetscLogGpuTimeEnd());

 80:   PetscCall(PetscLogGpuFlops(1.0 * N * N));
 81:   PetscCall(PetscDeviceContextSynchronizeIfWithBarrier_Internal(dctx));
 82:   PetscFunctionReturn(PETSC_SUCCESS);
 83: }

 85: #if PetscDefined(HAVE_CUDA)
 86: template struct UpperTriangular<DeviceType::CUDA>;
 87: #endif

 89: #if PetscDefined(HAVE_HIP)
 90: template struct UpperTriangular<DeviceType::HIP>;
 91: #endif

 93: } // namespace impl

 95: } // namespace cupm

 97: } // namespace device

 99: } // namespace Petsc

101: PETSC_INTERN PetscErrorCode MatUpperTriangularSolveInPlace_CUPM(PetscBool hermitian_transpose, PetscInt n, const PetscScalar A[], PetscInt lda, PetscScalar x[], PetscInt stride)
102: {
103:   using ::Petsc::device::cupm::impl::UpperTriangular;
104:   using ::Petsc::device::cupm::DeviceType;
105:   PetscDeviceContext dctx;
106:   PetscDeviceType    device_type;

108:   PetscFunctionBegin;
109:   PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
110:   PetscCall(PetscDeviceContextGetDeviceType(dctx, &device_type));
111:   switch (device_type) {
112: #if PetscDefined(HAVE_CUDA)
113:   case PETSC_DEVICE_CUDA:
114:     PetscCall(UpperTriangular<DeviceType::CUDA>::SolveInPlace(dctx, hermitian_transpose, n, A, lda, x, stride));
115:     break;
116: #endif
117: #if PetscDefined(HAVE_HIP)
118:   case PETSC_DEVICE_HIP:
119:     PetscCall(UpperTriangular<DeviceType::HIP>::SolveInPlace(dctx, hermitian_transpose, n, A, lda, x, stride));
120:     break;
121: #endif
122:   default:
123:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Unsupported device type %s", PetscDeviceTypes[device_type]);
124:   }
125:   PetscFunctionReturn(PETSC_SUCCESS);
126: }

128: PETSC_INTERN PetscErrorCode MatUpperTriangularSolveInPlaceCyclic_CUPM(PetscBool hermitian_transpose, PetscInt m, PetscInt oldest, PetscInt next, const PetscScalar A[], PetscInt lda, PetscScalar x[], PetscInt stride)
129: {
130:   using ::Petsc::device::cupm::impl::UpperTriangular;
131:   using ::Petsc::device::cupm::DeviceType;
132:   PetscDeviceContext dctx;
133:   PetscDeviceType    device_type;

135:   PetscFunctionBegin;
136:   PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
137:   PetscCall(PetscDeviceContextGetDeviceType(dctx, &device_type));
138:   switch (device_type) {
139: #if PetscDefined(HAVE_CUDA)
140:   case PETSC_DEVICE_CUDA:
141:     PetscCall(UpperTriangular<DeviceType::CUDA>::SolveInPlaceCyclic(dctx, hermitian_transpose, m, oldest, next, A, lda, x, stride));
142:     break;
143: #endif
144: #if PetscDefined(HAVE_HIP)
145:   case PETSC_DEVICE_HIP:
146:     PetscCall(UpperTriangular<DeviceType::HIP>::SolveInPlaceCyclic(dctx, hermitian_transpose, m, oldest, next, A, lda, x, stride));
147:     break;
148: #endif
149:   default:
150:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Unsupported device type %s", PetscDeviceTypes[device_type]);
151:   }
152:   PetscFunctionReturn(PETSC_SUCCESS);
153: }