[lang] Reorder sparse matrix before solving

python/taichi/examples/simulation/stable_fluid.py

@@ -27,7 +27,7 @@
                     help='The arch (backend) to run this example on')
 args, unknowns = parser.parse_known_args()
 
-res = 256
+res = 512
 dt = 0.03
 p_jacobi_iters = 500  # 40 for a quicker but less accurate result
 f_strength = 10000.0

taichi/program/sparse_solver.cpp

@@ -199,9 +199,11 @@ void CuSparseSolver::analyze_pattern(const SparseMatrix &sm) {
   SparseMatrix *sm_no_cv = const_cast<SparseMatrix *>(&sm);
   CuSparseMatrix *A = dynamic_cast<CuSparseMatrix *>(sm_no_cv);
   size_t rowsA = A->num_rows();
+  size_t colsA = A->num_cols();
   size_t nnzA = A->get_nnz();
   void *d_csrRowPtrA = A->get_row_ptr();
   void *d_csrColIndA = A->get_col_ind();
+  void *d_csrValA = A->get_val_ptr();
   CUSOLVERDriver::get_instance().csSpCreate(&cusolver_handle_);
   CUSPARSEDriver::get_instance().cpCreate(&cusparse_handel_);
   CUSPARSEDriver::get_instance().cpCreateMatDescr(&descr_);
@@ -210,12 +212,68 @@ void CuSparseSolver::analyze_pattern(const SparseMatrix &sm) {
   CUSPARSEDriver::get_instance().cpSetMatIndexBase(descr_,
                                                    CUSPARSE_INDEX_BASE_ZERO);
 
-  // step 1: create opaque info structure
+  // step 1: reorder the sparse matrix
+  float *h_csrValA = nullptr;
+  h_Q = (int *)malloc(sizeof(int) * colsA);
+  h_csrRowPtrB = (int *)malloc(sizeof(int) * (rowsA + 1));
+  h_csrColIndB = (int *)malloc(sizeof(int) * nnzA);
+  h_csrValB = (float *)malloc(sizeof(float) * nnzA);
+  h_csrValA = (float *)malloc(sizeof(float) * nnzA);
+  h_mapBfromA = (int *)malloc(sizeof(int) * nnzA);
+  assert(nullptr != h_Q);
+  assert(nullptr != h_csrRowPtrB);
+  assert(nullptr != h_csrColIndB);
+  assert(nullptr != h_csrValB);
+  assert(nullptr != h_mapBfromA);
+
+  CUDADriver::get_instance().memcpy_device_to_host(h_csrRowPtrB, d_csrRowPtrA,
+                                                   sizeof(int) * (rowsA + 1));
+  CUDADriver::get_instance().memcpy_device_to_host(h_csrColIndB, d_csrColIndA,
+                                                   sizeof(int) * nnzA);
+  CUDADriver::get_instance().memcpy_device_to_host(h_csrValA, d_csrValA,
+                                                   sizeof(float) * nnzA);
+
+  // compoute h_Q
+  CUSOLVERDriver::get_instance().csSpXcsrsymamdHost(
+      cusolver_handle_, rowsA, nnzA, descr_, h_csrRowPtrB, h_csrColIndB, h_Q);
+  CUDADriver::get_instance().malloc((void **)&d_Q, sizeof(int) * colsA);
+  CUDADriver::get_instance().memcpy_host_to_device((void *)d_Q, (void *)h_Q,
+                                                   sizeof(int) * (colsA));
+  size_t size_perm = 0;
+  CUSOLVERDriver::get_instance().csSpXcsrperm_bufferSizeHost(
+      cusolver_handle_, rowsA, colsA, nnzA, descr_, h_csrRowPtrB, h_csrColIndB,
+      h_Q, h_Q, &size_perm);
+  void *buffer_cpu = (void *)malloc(sizeof(char) * size_perm);
+  assert(nullptr != buffer_cpu);
+  for (int j = 0; j < nnzA; j++) {
+    h_mapBfromA[j] = j;
+  }
+  CUSOLVERDriver::get_instance().csSpXcsrpermHost(
+      cusolver_handle_, rowsA, colsA, nnzA, descr_, h_csrRowPtrB, h_csrColIndB,
+      h_Q, h_Q, h_mapBfromA, buffer_cpu);
+  // B = A( mapBfromA )
+  for (int j = 0; j < nnzA; j++) {
+    h_csrValB[j] = h_csrValA[h_mapBfromA[j]];
+  }
+  CUDADriver::get_instance().malloc((void **)&d_csrRowPtrB,
+                                    sizeof(int) * (rowsA + 1));
+  CUDADriver::get_instance().malloc((void **)&d_csrColIndB, sizeof(int) * nnzA);
+  CUDADriver::get_instance().malloc((void **)&d_csrValB, sizeof(float) * nnzA);
+  CUDADriver::get_instance().memcpy_host_to_device(
+      (void *)d_csrRowPtrB, (void *)h_csrRowPtrB, sizeof(int) * (rowsA + 1));
+  CUDADriver::get_instance().memcpy_host_to_device(
+      (void *)d_csrColIndB, (void *)h_csrColIndB, sizeof(int) * nnzA);
+  CUDADriver::get_instance().memcpy_host_to_device(
+      (void *)d_csrValB, (void *)h_csrValB, sizeof(float) * nnzA);
+  free(h_csrValA);
+  free(buffer_cpu);
+
+  // step 2: create opaque info structure
   CUSOLVERDriver::get_instance().csSpCreateCsrcholInfo(&info_);
 
-  // step 2: analyze chol(A) to know structure of L
+  // step 3: analyze chol(A) to know structure of L
   CUSOLVERDriver::get_instance().csSpXcsrcholAnalysis(
-      cusolver_handle_, rowsA, nnzA, descr_, d_csrRowPtrA, d_csrColIndA, info_);
+      cusolver_handle_, rowsA, nnzA, descr_, d_csrRowPtrB, d_csrColIndB, info_);
   is_analyzed_ = true;
 #else
   TI_NOT_IMPLEMENTED
@@ -229,26 +287,23 @@ void CuSparseSolver::factorize(const SparseMatrix &sm) {
   CuSparseMatrix *A = dynamic_cast<CuSparseMatrix *>(sm_no_cv);
   size_t rowsA = A->num_rows();
   size_t nnzA = A->get_nnz();
-  void *d_csrRowPtrA = A->get_row_ptr();
-  void *d_csrColIndA = A->get_col_ind();
-  void *d_csrValA = A->get_val_ptr();
 
   size_t size_internal = 0;
   size_t size_chol = 0;  // size of working space for csrlu
   // step 1: workspace for chol(A)
   CUSOLVERDriver::get_instance().csSpScsrcholBufferInfo(
-      cusolver_handle_, rowsA, nnzA, descr_, d_csrValA, d_csrRowPtrA,
-      d_csrColIndA, info_, &size_internal, &size_chol);
+      cusolver_handle_, rowsA, nnzA, descr_, d_csrValB, d_csrRowPtrB,
+      d_csrColIndB, info_, &size_internal, &size_chol);
 
   if (size_chol > 0)
     CUDADriver::get_instance().malloc(&gpu_buffer_, sizeof(char) * size_chol);
 
   // step 2: compute A = L*L^T
   CUSOLVERDriver::get_instance().csSpScsrcholFactor(
-      cusolver_handle_, rowsA, nnzA, descr_, d_csrValA, d_csrRowPtrA,
-      d_csrColIndA, info_, gpu_buffer_);
+      cusolver_handle_, rowsA, nnzA, descr_, d_csrValB, d_csrRowPtrB,
+      d_csrColIndB, info_, gpu_buffer_);
   // step 3: check if the matrix is singular
-  const double tol = 1.e-14;
+  const float tol = 1.e-14;
   int singularity = 0;
   CUSOLVERDriver::get_instance().csSpScsrcholZeroPivot(cusolver_handle_, info_,
                                                        tol, &singularity);
@@ -439,16 +494,46 @@ void CuSparseSolver::solve_rf(Program *prog,
   SparseMatrix *sm_no_cv = const_cast<SparseMatrix *>(&sm);
   CuSparseMatrix *A = dynamic_cast<CuSparseMatrix *>(sm_no_cv);
   size_t rowsA = A->num_rows();
+  size_t colsA = A->num_cols();
   size_t d_b = prog->get_ndarray_data_ptr_as_int(&b);
   size_t d_x = prog->get_ndarray_data_ptr_as_int(&x);
-  CUSOLVERDriver::get_instance().csSpScsrcholSolve(
-      cusolver_handle_, rowsA, (void *)d_b, (void *)d_x, info_, gpu_buffer_);
 
-  // TODO: free allocated memory and handles
-  // CUDADriver::get_instance().mem_free(gpu_buffer_);
-  // CUSOLVERDriver::get_instance().csSpDestory(cusolver_handle_);
-  // CUSPARSEDriver::get_instance().cpDestroy(cusparse_handel_);
-  // CUSPARSEDriver::get_instance().cpDestroyMatDescr(descrA);
+  // step 1: d_Qb = Q * b
+  void *d_Qb = nullptr;
+  CUDADriver::get_instance().malloc(&d_Qb, sizeof(float) * rowsA);
+  cusparseDnVecDescr_t vec_b;
+  cusparseSpVecDescr_t vec_Qb;
+  CUSPARSEDriver::get_instance().cpCreateDnVec(&vec_b, (int)rowsA, (void *)d_b,
+                                               CUDA_R_32F);
+  CUSPARSEDriver::get_instance().cpCreateSpVec(
+      &vec_Qb, (int)rowsA, (int)rowsA, d_Q, d_Qb, CUSPARSE_INDEX_32I,
+      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F);
+  CUSPARSEDriver::get_instance().cpGather(cusparse_handel_, vec_b, vec_Qb);
+
+  // step 2: solve B*z = Q*b using cholesky solver
+  void *d_z = nullptr;
+  CUDADriver::get_instance().malloc(&d_z, sizeof(float) * colsA);
+  CUSOLVERDriver::get_instance().csSpScsrcholSolve(
+      cusolver_handle_, rowsA, (void *)d_Qb, (void *)d_z, info_, gpu_buffer_);
+
+  // step 3: Q*x = z
+  cusparseSpVecDescr_t vecX;
+  cusparseDnVecDescr_t vecY;
+  CUSPARSEDriver::get_instance().cpCreateSpVec(
+      &vecX, (int)colsA, (int)colsA, d_Q, d_z, CUSPARSE_INDEX_32I,
+      CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F);
+  CUSPARSEDriver::get_instance().cpCreateDnVec(&vecY, (int)colsA, (void *)d_x,
+                                               CUDA_R_32F);
+  CUSPARSEDriver::get_instance().cpScatter(cusparse_handel_, vecX, vecY);
+
+  if (d_Qb != nullptr)
+    CUDADriver::get_instance().mem_free(d_Qb);
+  if (d_z != nullptr)
+    CUDADriver::get_instance().mem_free(d_z);
+  CUSPARSEDriver::get_instance().cpDestroySpVec(vec_Qb);
+  CUSPARSEDriver::get_instance().cpDestroyDnVec(vec_b);
+  CUSPARSEDriver::get_instance().cpDestroySpVec(vecX);
+  CUSPARSEDriver::get_instance().cpDestroyDnVec(vecY);
 #else
   TI_NOT_IMPLEMENTED
 #endif
@@ -492,6 +577,38 @@ std::unique_ptr<SparseSolver> make_sparse_solver(DataType dt,
     TI_ERROR("Not supported sparse solver type: {}", solver_type);
 }
 
+CuSparseSolver::~CuSparseSolver() {
+#if defined(TI_WITH_CUDA)
+  if (h_Q != nullptr)
+    free(h_Q);
+  if (h_csrRowPtrB != nullptr)
+    free(h_csrRowPtrB);
+  if (h_csrColIndB != nullptr)
+    free(h_csrColIndB);
+  if (h_csrValB != nullptr)
+    free(h_csrValB);
+  if (h_mapBfromA != nullptr)
+    free(h_mapBfromA);
+  if (info_ != nullptr)
+    CUSOLVERDriver::get_instance().csSpDestroyCsrcholInfo(info_);
+  if (cusolver_handle_ != nullptr)
+    CUSOLVERDriver::get_instance().csSpDestory(cusolver_handle_);
+  if (cusparse_handel_ != nullptr)
+    CUSPARSEDriver::get_instance().cpDestroy(cusparse_handel_);
+  if (descr_ != nullptr)
+    CUSPARSEDriver::get_instance().cpDestroyMatDescr(descr_);
+  if (gpu_buffer_ != nullptr)
+    CUDADriver::get_instance().mem_free(gpu_buffer_);
+  if (d_Q != nullptr)
+    CUDADriver::get_instance().mem_free(d_Q);
+  if (d_csrRowPtrB != nullptr)
+    CUDADriver::get_instance().mem_free(d_csrRowPtrB);
+  if (d_csrColIndB != nullptr)
+    CUDADriver::get_instance().mem_free(d_csrColIndB);
+  if (d_csrValB != nullptr)
+    CUDADriver::get_instance().mem_free(d_csrValB);
+#endif
+}
 std::unique_ptr<SparseSolver> make_cusparse_solver(
     DataType dt,
     const std::string &solver_type,

taichi/program/sparse_solver.h

@@ -85,9 +85,18 @@ class CuSparseSolver : public SparseSolver {
   bool is_analyzed_{false};
   bool is_factorized_{false};
 
+  int *h_Q{nullptr}; /* <int> n,  B = Q*A*Q' or B = A(Q,Q) by MATLAB notation */
+  int *d_Q{nullptr};
+  int *h_csrRowPtrB{nullptr}; /* <int> n+1 */
+  int *h_csrColIndB{nullptr}; /* <int> nnzA */
+  float *h_csrValB{nullptr};  /* <float> nnzA */
+  int *h_mapBfromA{nullptr};  /* <int> nnzA */
+  int *d_csrRowPtrB{nullptr}; /* <int> n+1 */
+  int *d_csrColIndB{nullptr}; /* <int> nnzA */
+  float *d_csrValB{nullptr};  /* <float> nnzA */
  public:
   CuSparseSolver();
-  ~CuSparseSolver() override = default;
+  ~CuSparseSolver() override;
   bool compute(const SparseMatrix &sm) override {
     TI_NOT_IMPLEMENTED;
   };

taichi/rhi/cuda/cusolver_functions.inc.h

@@ -9,12 +9,14 @@ PER_CUSOLVER_FUNCTION(csSpCreate, cusolverSpCreate, cusolverSpHandle_t * );
 PER_CUSOLVER_FUNCTION(csSpDestory, cusolverSpDestroy, cusolverSpHandle_t );
 PER_CUSOLVER_FUNCTION(csSpXcsrissymHost, cusolverSpXcsrissymHost, cusolverSpHandle_t,int ,int ,const cusparseMatDescr_t ,const void *,const void *,const void *,void *);
 PER_CUSOLVER_FUNCTION(csSpXcsrsymrcmHost, cusolverSpXcsrsymrcmHost, cusolverSpHandle_t ,int ,int ,const cusparseMatDescr_t ,const void *,const void *,void *);
+PER_CUSOLVER_FUNCTION(csSpXcsrsymamdHost, cusolverSpXcsrsymamdHost, cusolverSpHandle_t ,int ,int ,const cusparseMatDescr_t ,const void *,const void *,void *);
 PER_CUSOLVER_FUNCTION(csSpXcsrperm_bufferSizeHost, cusolverSpXcsrperm_bufferSizeHost, cusolverSpHandle_t ,int ,int ,int ,const cusparseMatDescr_t ,void *,void *,const void *,const void *,size_t *);
 PER_CUSOLVER_FUNCTION(csSpXcsrpermHost, cusolverSpXcsrpermHost, cusolverSpHandle_t ,int ,int ,int ,const cusparseMatDescr_t ,void *,void *,const void *,const void *,void *,void *);
 PER_CUSOLVER_FUNCTION(csSpScsrlsvcholHost, cusolverSpScsrlsvcholHost, cusolverSpHandle_t ,int ,int ,const cusparseMatDescr_t ,const void *,const void *,const void *,const void *,float ,int ,void *,void *);
 
 // cusolver preview API
 PER_CUSOLVER_FUNCTION(csSpCreateCsrcholInfo, cusolverSpCreateCsrcholInfo, csrcholInfo_t*);
+PER_CUSOLVER_FUNCTION(csSpDestroyCsrcholInfo, cusolverSpDestroyCsrcholInfo, csrcholInfo_t);
 PER_CUSOLVER_FUNCTION(csSpXcsrcholAnalysis, cusolverSpXcsrcholAnalysis,  cusolverSpHandle_t,int ,int ,const cusparseMatDescr_t , void *, void *,csrcholInfo_t );
 PER_CUSOLVER_FUNCTION(csSpScsrcholBufferInfo, cusolverSpScsrcholBufferInfo,  cusolverSpHandle_t ,int ,int ,const cusparseMatDescr_t ,void *,void *,void *,csrcholInfo_t ,size_t *,size_t *);
 PER_CUSOLVER_FUNCTION(csSpScsrcholFactor, cusolverSpScsrcholFactor, cusolverSpHandle_t ,int ,int ,const cusparseMatDescr_t ,void *,void *,void *,csrcholInfo_t ,void *);

taichi/rhi/cuda/cusparse_functions.inc.h

@@ -20,6 +20,8 @@ PER_CUSPARSE_FUNCTION(cpCreateIdentityPermutation, cusparseCreateIdentityPermuta
 PER_CUSPARSE_FUNCTION(cpXcoosort_bufferSizeExt, cusparseXcoosort_bufferSizeExt, cusparseHandle_t,int ,int,int, void* ,void* ,void*);
 PER_CUSPARSE_FUNCTION(cpXcoosortByRow, cusparseXcoosortByRow, cusparseHandle_t,int,int,int,void* ,void* ,void* ,void*);
 PER_CUSPARSE_FUNCTION(cpGather, cusparseGather, cusparseHandle_t, cusparseDnVecDescr_t, cusparseSpVecDescr_t);
+PER_CUSPARSE_FUNCTION(cpScatter, cusparseScatter, cusparseHandle_t, cusparseSpVecDescr_t, cusparseDnVecDescr_t);
+PER_CUSPARSE_FUNCTION(cpSsctr, cusparseSsctr, cusparseHandle_t, int, void*, void*, void*, cusparseIndexBase_t);
 PER_CUSPARSE_FUNCTION(cpSetPointerMode, cusparseSetPointerMode, cusparseHandle_t, cusparsePointerMode_t);
 PER_CUSPARSE_FUNCTION(cpCsrSetPointers, cusparseCsrSetPointers, cusparseSpMatDescr_t,void*, void*, void*);