diff --git a/legion/spmv/main.cpp b/legion/spmv/main.cpp index 02f9d0651..10f8ebebe 100644 --- a/legion/spmv/main.cpp +++ b/legion/spmv/main.cpp @@ -228,7 +228,7 @@ 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) { @@ -236,7 +236,9 @@ void convertDISTALCSRToStandardCSR(Context ctx, Runtime* runtime, int pieces, Le 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& 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,9 +370,9 @@ void top_level_task(const Task* task, const std::vector& 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. @@ -361,7 +380,7 @@ void top_level_task(const Task* task, const std::vector& regions 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); diff --git a/legion/spmv/taco-generated.cpp b/legion/spmv/taco-generated.cpp index 84216825b..9160e882a 100644 --- a/legion/spmv/taco-generated.cpp +++ b/legion/spmv/taco-generated.cpp @@ -121,6 +121,9 @@ void task_1(const Task* task, const std::vector& regions, Contex auto B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); auto B2_crd_accessor = createAccessor(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& 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(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(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(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(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(registrar, "task_5"); } diff --git a/legion/spmv/taco-generated.cu b/legion/spmv/taco-generated.cu index 0090eb667..c7627fdc6 100644 --- a/legion/spmv/taco-generated.cu +++ b/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& 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& 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(B_vals, FID_VAL); auto c_vals_ro_accessor = createAccessor(c_vals, FID_VAL); auto a_vals_rw_accessor = createAccessor(a_vals, FID_VAL); - auto B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); + auto B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); auto B2_crd_accessor = createAccessor(B2_crd, FID_COORD); auto posDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_pos)); @@ -131,20 +147,25 @@ void task_1(const Task* task, const std::vector& 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 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<<>>(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), @@ -152,7 +173,7 @@ void task_1(const Task* task, const std::vector& regions, Contex 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,7 +193,7 @@ void task_1(const Task* task, const std::vector& regions, Contex handle, CUSPARSE_ALG_MERGE_PATH, CUSPARSE_OPERATION_NON_TRANSPOSE, - (posDomain.get_volume() - 1), + posDomain.get_volume(), B2_dimension, crdDomain.get_volume(), &(alpha), @@ -180,7 +201,7 @@ void task_1(const Task* task, const std::vector& regions, Contex 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& 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) { diff --git a/legion/spmv/taco-generated.cuh b/legion/spmv/taco-generated.cuh index d16bfb661..71afe9f66 100644 --- a/legion/spmv/taco-generated.cuh +++ b/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); diff --git a/legion/spmv/taco-generated.h b/legion/spmv/taco-generated.h index 0eec988ff..ea3b0d8a2 100644 --- a/legion/spmv/taco-generated.h +++ b/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); diff --git a/legion/spmv/taco-header.h b/legion/spmv/taco-header.h new file mode 100644 index 000000000..168f67961 --- /dev/null +++ b/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 diff --git a/petsc/benchmark.cpp b/petsc/benchmark.cpp index 618994769..12256ec05 100644 --- a/petsc/benchmark.cpp +++ b/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);