legion/spmv: several bugfixes for running SpMV weak scale

legion/spmv/main.cpp

@@ -228,15 +228,17 @@ void convertDISTALCSRToStandardCSR(Context ctx, Runtime* runtime, int pieces, Le
 
   // Delete the old region, as this is important to not OOM on some benchmarks.
   // Before we do so though, detach the region from the external allocation.
-  size_t toDelete = 0;
+  int toDelete = -1;
   for (size_t i = 0; i < attached.attachedRegions.size(); i++) {
     auto phys = attached.attachedRegions[i];
     if (phys.get_logical_region() == oldReg) {
       toDelete = i;
       runtime->detach_external_resource(ctx, phys);
     }
   }
-  attached.attachedRegions.erase(attached.attachedRegions.begin() + toDelete);
+  if (!attached.attachedRegions.empty() && toDelete != -1) {
+    attached.attachedRegions.erase(attached.attachedRegions.begin() + toDelete);
+  }
   runtime->destroy_logical_region(ctx, oldReg);
 }
 
@@ -325,7 +327,24 @@ void top_level_task(const Task* task, const std::vector<PhysicalRegion>& regions
       launcher.add_region_requirement(RegionRequirement(xEqLPart, 0, WRITE_ONLY, EXCLUSIVE, x.valsParent).add_field(FID_VAL));
       runtime->execute_index_space(ctx, launcher);
     }
-  }
+  } 
+
+  LogicalRegion replX;
+  LogicalPartition replXLPart;
+  if (weakScale) {
+    // To get around not having collective instances we need to manually replicate the x vector.
+    auto replIndexSpace = runtime->create_index_space(ctx, Rect<1>(0, (size_t(pieces) * size_t(A.dims[1])) - 1));
+    replX = runtime->create_logical_region(ctx, replIndexSpace, y.vals.get_field_space());
+    DomainPointColoring xColoring;
+    for (size_t i = 0; i < pieces; i++) {
+      size_t lo = i * size_t(A.dims[1]);
+      size_t hi = ((i + 1) * size_t(A.dims[1])) - 1;
+      xColoring[i] = Rect<1>(lo, hi);
+    }
+    auto xIndexPart = runtime->create_index_partition(ctx, replIndexSpace, eqPartDomain, xColoring, LEGION_DISJOINT_COMPLETE_KIND);
+    replXLPart = runtime->get_logical_partition(ctx, replX, xIndexPart);
+    runtime->fill_field(ctx, replX, replX, FID_VAL, valType(1));
+  } 
 
   // Create a partition of y for forcing reductions.
   auto yEqIndexPart = runtime->create_equal_partition(ctx, y.vals.get_index_space(), eqPartIspace);
@@ -351,17 +370,17 @@ void top_level_task(const Task* task, const std::vector<PhysicalRegion>& regions
   // If we are using CUDA and using a row-split schedule, we're going to use CuSparse
   // at the leaves. Therefore, we need to convert the DISTAL pos region into a standard
   // CSR region to be compatible with CuSparse.
-  if (conf == CSR_ROW_SPLIT) {
-    convertDISTALCSRToStandardCSR(ctx, runtime, pieces, A, rowPack, Aex);
-  }
+  // if (conf == CSR_ROW_SPLIT) {
+  //   convertDISTALCSRToStandardCSR(ctx, runtime, pieces, A, rowPack, Aex);
+  // }
 #endif
 
   // Benchmark the computation.
   auto avgTime = benchmarkAsyncCallWithWarmup(ctx, runtime, warmup, n, [&]() {
     if (dump) { runtime->fill_field(ctx, y.vals, y.valsParent, FID_VAL, valType(0)); }
     switch (conf) {
       case CSR_ROW_SPLIT:
-        computeLegionRowSplit(ctx, runtime, &y, &A, &x, &rowPack, pieces); break;
+        computeLegionRowSplit(ctx, runtime, &y, &A, &x, &rowPack, pieces, replX, replXLPart); break;
       case CSR_POS_SPLIT:
         computeLegionPosSplit(ctx, runtime, &y, &A, &x, &posPack, pieces);
         launchDummyReadOverPartition(ctx, runtime, y.vals, yEqLPart, FID_VAL, eqPartDomain);

legion/spmv/taco-generated.cpp

@@ -121,6 +121,9 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
   auto B2_pos_accessor = createAccessor<AccessorRORect_1_1>(B2_pos, FID_RECT_1);
   auto B2_crd_accessor = createAccessor<AccessorROint32_t1>(B2_crd, FID_COORD);
 
+  // Have to use a raw pointer for manual access here.
+  auto c_vals_ptr = c_vals_ro_accessor.ptr(runtime->get_index_space_domain(ctx, c_vals.get_logical_region().get_index_space()).lo());
+
   int64_t pointID1 = io;
   #pragma omp parallel for schedule(dynamic, 128)
   for (int64_t ii = 0; ii < ((B1_dimension + (pieces - 1)) / pieces); ii++) {
@@ -137,12 +140,12 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
     for (int64_t jB = B2_pos_accessor[Point<1>(i)].lo; jB < (B2_pos_accessor[Point<1>(i)].hi + 1); jB++) {
       int64_t j = B2_crd_accessor[(jB * 1)];
       int64_t jc = j;
-      a_vals_rw_accessor[Point<1>(i)] = a_vals_rw_accessor[Point<1>(i)] + B_vals_ro_accessor[Point<1>(jB)] * c_vals_ro_accessor[Point<1>(j)];
+      a_vals_rw_accessor[Point<1>(i)] = a_vals_rw_accessor[Point<1>(i)] + B_vals_ro_accessor[Point<1>(jB)] * c_vals_ptr[j];
     }
   }
 }
 
-void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces) {
+void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces, LogicalRegion replX, LogicalPartition replXLPart) {
   auto a_vals_parent = a->valsParent;
   int B1_dimension = B->dims[0];
   RegionWrapper B2_crd = B->indices[1][1];
@@ -167,7 +170,8 @@ void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, Legion
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.indicesPartitions[1][0], 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1));
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.indicesPartitions[1][1], 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD));
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.valsPartition, 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL));
-  launcher.add_region_requirement(RegionRequirement(get_logical_region(c_vals), READ_ONLY, EXCLUSIVE, c_vals_parent).add_field(FID_VAL));
+  launcher.add_region_requirement(RegionRequirement(replXLPart, 0, READ_ONLY, EXCLUSIVE, replX).add_field(FID_VAL));
+  // launcher.add_region_requirement(RegionRequirement(get_logical_region(c_vals), READ_ONLY, EXCLUSIVE, c_vals_parent).add_field(FID_VAL));
   runtime->execute_index_space(ctx, launcher);
 
 }
@@ -203,6 +207,7 @@ partitionPackForcomputeLegionPosSplit partitionForcomputeLegionPosSplit(Legion::
   LogicalPartition B2_crd_part = runtime->get_logical_partition(ctx, B2_crd, B2_crd_index_part);
   IndexPartition posSparsePartB2 = runtime->create_partition_by_preimage_range(ctx, B2_crd_index_part, B2_pos, B2_pos_parent, FID_RECT_1, runtime->get_index_partition_color_space_name(ctx, B2_crd_index_part));
   IndexPartition posIndexPartB2 = densifyPartition(ctx, runtime, get_index_space(B2_pos), posSparsePartB2);
+  runtime->destroy_index_partition(ctx, posSparsePartB2, false /* unordered */, true /* recurse */);
   LogicalPartition posPartB2 = runtime->get_logical_partition(ctx, B2_pos, posIndexPartB2);
   LogicalPartition BValsLogicalPart = copyPartition(ctx, runtime, B2_crd_part, B_vals);
   IndexPartition BDenseRun0Partition = copyPartition(ctx, runtime, posPartB2, B_dense_run_0);
@@ -748,31 +753,31 @@ void computeLegionCSCMSpV(Legion::Context ctx, Legion::Runtime* runtime, LegionT
 void registerTacoTasks() {
   {
     TaskVariantRegistrar registrar(taskID(1), "task_1");
-    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC));
     registrar.set_leaf();
     Runtime::preregister_task_variant<task_1>(registrar, "task_1");
   }
   {
     TaskVariantRegistrar registrar(taskID(2), "task_2");
-    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC));
     registrar.set_leaf();
     Runtime::preregister_task_variant<task_2>(registrar, "task_2");
   }
   {
     TaskVariantRegistrar registrar(taskID(3), "task_3");
-    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC));
     registrar.set_leaf();
     Runtime::preregister_task_variant<task_3>(registrar, "task_3");
   }
   {
     TaskVariantRegistrar registrar(taskID(4), "task_4");
-    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC));
     registrar.set_leaf();
     Runtime::preregister_task_variant<task_4>(registrar, "task_4");
   }
   {
     TaskVariantRegistrar registrar(taskID(5), "task_5");
-    registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
+    registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC));
     registrar.set_leaf();
     Runtime::preregister_task_variant<task_5>(registrar, "task_5");
   }

legion/spmv/taco-generated.cu

@@ -102,6 +102,18 @@ partitionPackForcomputeLegionRowSplit partitionForcomputeLegionRowSplit(Legion::
   return computePartitions;
 }
 
+__global__
+void constructValidPos(size_t elems, const Rect<1>* posGlobal, int32_t* posLocal) {
+  const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= elems) {
+    return;
+  }
+  posLocal[idx] = posGlobal[idx].lo - posGlobal[0].lo;
+  if (idx == 0) {
+    posLocal[elems] = posGlobal[elems - 1].hi + 1;
+  }
+}
+
 void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Context ctx, Runtime* runtime) {
   PhysicalRegion a_vals = regions[0];
   LogicalRegion a_vals_parent = regions[0].get_logical_region();
@@ -119,10 +131,14 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
   int64_t B2_dimension = args->B2_dimension;
   int32_t pieces = args->pieces;
 
+  cusparseHandle_t handle = getCuSparse();
+  cudaStream_t taskStream = cudaStream_t();
+  cudaStreamCreate(&(taskStream));
+
   auto B_vals_ro_accessor = createAccessor<AccessorROdouble1>(B_vals, FID_VAL);
   auto c_vals_ro_accessor = createAccessor<AccessorROdouble1>(c_vals, FID_VAL);
   auto a_vals_rw_accessor = createAccessor<AccessorRWdouble1>(a_vals, FID_VAL);
-  auto B2_pos_accessor = createAccessor<AccessorROint32_t1>(B2_pos, FID_RECT_1);
+  auto B2_pos_accessor = createAccessor<AccessorRORect_1_1>(B2_pos, FID_RECT_1);
   auto B2_crd_accessor = createAccessor<AccessorROint32_t1>(B2_crd, FID_COORD);
 
   auto posDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_pos));
@@ -131,28 +147,33 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
   auto BValsDomain = runtime->get_index_space_domain(ctx, get_index_space(B_vals));
   auto cValsDomain = runtime->get_index_space_domain(ctx, get_index_space(c_vals));
   double alpha = 1.0000000000000000;
-  cusparseHandle_t handle = getCuSparse();
-  cudaStream_t taskStream = cudaStream_t();
-  cudaStreamCreate(&(taskStream));
   CHECK_CUSPARSE(cusparseSetStream(handle, taskStream));
   cusparseMatDescr_t mat = cusparseMatDescr_t();
   CHECK_CUSPARSE(cusparseCreateMatDescr(&(mat)));
   CHECK_CUSPARSE(cusparseSetMatIndexBase(mat, CUSPARSE_INDEX_BASE_ZERO));
   CHECK_CUSPARSE(cusparseSetMatType(mat, CUSPARSE_MATRIX_TYPE_GENERAL));
+
+  // We have to convert this globally addressed pos array into a local pos array.
   uint64_t bufferSize = 0;
+  Legion::DeferredBuffer<int32_t, 1> posBuf(Rect<1>(0, posDomain.get_volume()), Memory::Kind::GPU_FB_MEM);
+  const int THREADS_PER_BLOCK = 256;
+  auto blocks = (posDomain.get_volume() + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+  taco_iassert(blocks != 0);
+  constructValidPos<<<blocks, THREADS_PER_BLOCK, 0, taskStream>>>(posDomain.get_volume(), B2_pos_accessor.ptr(posDomain.lo()), posBuf.ptr(0));
+
   CHECK_CUSPARSE(cusparseCsrmvEx_bufferSize(
     handle,
     CUSPARSE_ALG_MERGE_PATH,
     CUSPARSE_OPERATION_NON_TRANSPOSE,
-    (posDomain.get_volume() - 1),
+    posDomain.get_volume(),
     B2_dimension,
     crdDomain.get_volume(),
     &(alpha),
     CUDA_R_64F,
     mat,
     B_vals_ro_accessor.ptr(BValsDomain.lo()),
     CUDA_R_64F,
-    B2_pos_accessor.ptr(posDomain.lo()),
+    posBuf.ptr(0),
     B2_crd_accessor.ptr(crdDomain.lo()),
     c_vals_ro_accessor.ptr(cValsDomain.lo()),
     CUDA_R_64F,
@@ -172,15 +193,15 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
     handle,
     CUSPARSE_ALG_MERGE_PATH,
     CUSPARSE_OPERATION_NON_TRANSPOSE,
-    (posDomain.get_volume() - 1),
+    posDomain.get_volume(),
     B2_dimension,
     crdDomain.get_volume(),
     &(alpha),
     CUDA_R_64F,
     mat,
     B_vals_ro_accessor.ptr(BValsDomain.lo()),
     CUDA_R_64F,
-    B2_pos_accessor.ptr(posDomain.lo()),
+    posBuf.ptr(0),
     B2_crd_accessor.ptr(crdDomain.lo()),
     c_vals_ro_accessor.ptr(cValsDomain.lo()),
     CUDA_R_64F,
@@ -193,7 +214,7 @@ void task_1(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
   ));
 }
 
-void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces) {
+void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces, LogicalRegion replX, LogicalPartition replXLPart) {
   auto a_vals_parent = a->valsParent;
   int B2_dimension = B->dims[1];
   auto B2_pos_parent = B->indicesParents[1][0];
@@ -216,9 +237,9 @@ void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, Legion
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.indicesPartitions[1][0], 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1));
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.indicesPartitions[1][1], 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD));
   launcher.add_region_requirement(RegionRequirement(partitionPack->BPartition.valsPartition, 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL));
-  launcher.add_region_requirement(RegionRequirement(get_logical_region(c_vals), READ_ONLY, EXCLUSIVE, c_vals_parent).add_field(FID_VAL));
+  launcher.add_region_requirement(RegionRequirement(replXLPart, 0, READ_ONLY, EXCLUSIVE, replX).add_field(FID_VAL));
+  // launcher.add_region_requirement(RegionRequirement(get_logical_region(c_vals), READ_ONLY, EXCLUSIVE, c_vals_parent).add_field(FID_VAL));
   runtime->execute_index_space(ctx, launcher);
-
 }
 
 partitionPackForcomputeLegionPosSplit partitionForcomputeLegionPosSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, int32_t pieces) {

legion/spmv/taco-generated.cuh

@@ -26,7 +26,7 @@ struct partitionPackForcomputeLegionCSCMSpV {
 partitionPackForcomputeLegionRowSplit partitionForcomputeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, int32_t pieces);
 
 
-void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces);
+void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces, Legion::LogicalRegion replX, Legion::LogicalPartition replXLPart);
 
 partitionPackForcomputeLegionPosSplit partitionForcomputeLegionPosSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, int32_t pieces);
 

legion/spmv/taco-generated.h

@@ -31,7 +31,7 @@ struct partitionPackForcomputeLegionCSCMSpV {
 partitionPackForcomputeLegionRowSplit partitionForcomputeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, int32_t pieces);
 
 
-void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces);
+void computeLegionRowSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, partitionPackForcomputeLegionRowSplit* partitionPack, int32_t pieces, Legion::LogicalRegion replX, Legion::LogicalPartition replXLPart);
 
 partitionPackForcomputeLegionPosSplit partitionForcomputeLegionPosSplit(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* a, LegionTensor* B, LegionTensor* c, int32_t pieces);
 

legion/spmv/taco-header.h

@@ -0,0 +1,14 @@
+#ifndef TACO_TEMP_H
+#define TACO_TEMP_H
+
+#include "legion.h"
+
+struct partitionPackForcomputeLegionRowSplit {
+  LegionTensorPartition aPartition;
+  LegionTensorPartition BPartition;
+};
+
+// typedef Legion::Rect<1, int32_t> PosRect;
+typedef Legion::Rect<1> PosRect;
+
+#endif

petsc/benchmark.cpp

@@ -114,11 +114,15 @@ void spmvWeakScale(int warmup, int niter) {
   const size_t NNZ_PER_ROW = 501;
   PetscInt size;
   MPI_Comm_size(PETSC_COMM_WORLD, &size);
-  // TODO (rohany): Extract this logic to make it composable?
-  // Get the number of processes per node.
-  size = size / 40;
+  // Get the number of processes per node. If we have GPUs then there's
+  // just 1 per GPU.
+  if (!GPUENABLED) {
+    size = size / 40;
+  }
   int dim = NNZ_TARGET / NNZ_PER_ROW * size_t(size);
-
+
+  // I get segfaults when I directly make this an MPIAIJ matrix, so we require passing
+  // in `-mat_type mpiaij` in the command line for this to work correctly for GPUs. 
   Mat B;
   MatCreate(PETSC_COMM_WORLD, &B);
   MatSetSizes(B, PETSC_DECIDE, PETSC_DECIDE, dim, dim);
@@ -128,7 +132,6 @@ void spmvWeakScale(int warmup, int niter) {
 
   // Fill the bands of B up.
   PetscInt rStart, rEnd;
-  // TODO (rohany): This call requires the mpiaij matrix type.
   MatGetOwnershipRange(B, &rStart, &rEnd);
   for (int i = rStart; i < rEnd; i++) {
     for (int j = (i - (int(NNZ_PER_ROW) / 2)); j <= (i + (int(NNZ_PER_ROW) / 2)); j++) {
@@ -151,8 +154,17 @@ void spmvWeakScale(int warmup, int niter) {
   VecSet(x, one);
   VecSet(y, zero);
 
+  // Insertion seems to be too slow into a MPIAIJCUSPARSE matrix by default, so we
+  // have to make a normal MPIAIJ matrix and then convert it into MPIAIJCUSPARSE.
+  Mat BFinal;
+  if (GPUENABLED) {
+    MatConvert(B, MATMPIAIJCUSPARSE, MAT_INITIAL_MATRIX, &BFinal);
+  } else {
+    BFinal = B;
+  }
+
   auto avgTime = benchmarkWithWarmup(warmup, niter, [&]() {
-    MatMult(B, x, y);
+    MatMult(BFinal, x, y);
   }); 
 
   PetscPrintf(PETSC_COMM_WORLD, "Average time: %lf ms.\n", avgTime * 1000);