diff --git a/legion/include/taco_mapper.h b/legion/include/taco_mapper.h index ad50a8b0d..a0f455dce 100644 --- a/legion/include/taco_mapper.h +++ b/legion/include/taco_mapper.h @@ -197,7 +197,9 @@ class TACOMapper : public Legion::Mapping::DefaultMapper { // InFlightTask represents a task currently being executed. struct InFlightTask { // Unique identifier of the task instance. - Legion::UniqueID id; + // Legion::UniqueID id; + // TODO (rohany): Comment. + std::pair id; // An event that will be triggered when the task finishes. Legion::Mapping::MapperEvent event; // A clock measurement from when the task was scheduled. diff --git a/legion/spinnerprod/taco-generated.cpp b/legion/spinnerprod/taco-generated.cpp index 1fc2ce963..a2be009b3 100644 --- a/legion/spinnerprod/taco-generated.cpp +++ b/legion/spinnerprod/taco-generated.cpp @@ -458,13 +458,13 @@ double computeLegionDDS(Legion::Context ctx, Legion::Runtime* runtime, LegionTen void registerTacoTasks() { { TaskVariantRegistrar registrar(taskID(1), "task_1"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_1"); } { TaskVariantRegistrar registrar(taskID(2), "task_2"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_2"); } diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 75a954004..d169c8415 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -6,17 +6,21 @@ #include "legion_utils.h" #include "legion_string_utils.h" #include "error.h" +#ifdef TACO_USE_CUDA +#include "taco-generated.cuh" +#else #include "taco-generated.h" +#endif using namespace Legion; typedef double valType; void top_level_task(const Task* task, const std::vector& regions, Context ctx, Runtime* runtime) { std::string csrFileName; - bool dump = false; + bool dump = false, consMem = false; // The j-dimension if the computation will commonly have a small value // that is divisible by 32, as per Stephen and Chang-wan. - int n = 10, pieces = 0, warmup = 5, jDim = 32; + int n = 10, pieces = 0, warmup = 5, jDim = 32, gx = 0, gy = 0; Realm::CommandLineParser parser; parser.add_option_string("-tensor", csrFileName); parser.add_option_bool("-dump", dump); @@ -24,12 +28,16 @@ void top_level_task(const Task* task, const std::vector& regions parser.add_option_int("-pieces", pieces); parser.add_option_int("-warmup", warmup); parser.add_option_int("-jdim", jDim); + // Arguments for memory-conserving schedule. + parser.add_option_bool("-consMem", consMem); + parser.add_option_int("-gx", gx); + parser.add_option_int("-gy", gy); auto args = Runtime::get_input_args(); taco_uassert(parser.parse_command_line(args.argc, args.argv)) << "Parse failure."; taco_uassert(!csrFileName.empty()) << "Provide a matrix with -tensor"; // Figure out how many pieces to chop up the data into. - if (pieces == 0) { + if (!consMem && pieces == 0) { pieces = getNumPieces(ctx, runtime); taco_uassert(pieces != 0) << "Please provide a number of pieces to split into with -pieces. Unable to automatically find."; } @@ -41,20 +49,41 @@ void top_level_task(const Task* task, const std::vector& regions runtime->fill_field(ctx, A.vals, A.valsParent, FID_VAL, valType(0)); runtime->fill_field(ctx, C.vals, C.valsParent, FID_VAL, valType(1)); - auto pack = partitionForcomputeLegion(ctx, runtime, &A, &B, &C, pieces); +// #ifdef TACO_USE_CUDA + partitionPackForcomputeLegion pack; + partitionPackForcomputeLegionConsMem packCons; + if (consMem) { + packCons = partitionForcomputeLegionConsMem(ctx, runtime, &A, &B, &C, gx, gy); + } else { + pack = partitionForcomputeLegion(ctx, runtime, &A, &B, &C, pieces); + } +// #else +// auto pack = partitionForcomputeLegion(ctx, runtime, &A, &B, &C, pieces); +// taco_iassert(!consMem); +// #endif - auto commPart = createSparseAliasingPartitions(ctx, runtime, A.vals.get_index_space(), pack.APartition.valsPartition.get_index_partition()); - auto commLPart = runtime->get_logical_partition(ctx, A.vals, commPart); + + LogicalPartition commLPart; + if (!consMem) { + auto commPart = createSparseAliasingPartitions(ctx, runtime, A.vals.get_index_space(), pack.APartition.valsPartition.get_index_partition()); + commLPart = runtime->get_logical_partition(ctx, A.vals, commPart); + } auto avgTime = benchmarkAsyncCallWithWarmup(ctx, runtime, warmup, n, [&]() { if (dump) { runtime->fill_field(ctx, A.vals, A.valsParent, FID_VAL, valType(0)); } - computeLegion(ctx, runtime, &A, &B, &C, &pack, pieces); -#ifdef TACO_USE_CUDA - // Collapse our reduction buffers. We use sparse instances to force just the communication - // that we want. We only do this for the GPU schedule, as the CPU schedule does not - // use Legion reductions. - launchDummyReadOverPartition(ctx, runtime, A.vals, commLPart, FID_VAL, Rect<1>(0, pieces - 1), false /* wait */, true /* untrack */, false /* cpuOnly */, true /* sparse */); -#endif +// #ifdef TACO_USE_CUDA + if (consMem) { + computeLegionConsMem(ctx, runtime, &A, &B, &C, &packCons, gx, gy); + } else { + computeLegion(ctx, runtime, &A, &B, &C, &pack, pieces); + // Collapse our reduction buffers. We use sparse instances to force just the communication + // that we want. We only do this for the GPU schedule, as the CPU schedule does not + // use Legion reductions. + launchDummyReadOverPartition(ctx, runtime, A.vals, commLPart, FID_VAL, Rect<1>(0, pieces - 1), false /* wait */, true /* untrack */, false /* cpuOnly */, true /* sparse */); + } +// #else +// computeLegion(ctx, runtime, &A, &B, &C, &pack, pieces); +// #endif }); LEGION_PRINT_ONCE(runtime, ctx, stdout, "Average execution time: %lf ms\n", avgTime); diff --git a/legion/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index 996552872..f92155d10 100644 --- a/legion/spmm/taco-generated.cpp +++ b/legion/spmm/taco-generated.cpp @@ -18,6 +18,14 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t A2_dimension; + int64_t B1_dimension; + int64_t C2_dimension; + int32_t gx; + int32_t gy; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; @@ -164,12 +172,192 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* 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); +} + +partitionPackForcomputeLegionConsMem partitionForcomputeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx, int32_t gy) { + RegionWrapper A_vals = A->vals; + IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; + int B1_dimension = B->dims[0]; + int B2_dimension = B->dims[1]; + RegionWrapper B2_pos = B->indices[1][0]; + RegionWrapper B2_crd = B->indices[1][1]; + auto B2_pos_parent = B->indicesParents[1][0]; + RegionWrapper B_vals = B->vals; + IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; + int C2_dimension = C->dims[1]; + RegionWrapper C_vals = C->vals; + IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + Point<2> lowerBound = Point<2>(0, 0); + Point<2> upperBound = Point<2>((gx - 1), (gy - 1)); + auto distFusedIndexSpace = runtime->create_index_space(ctx, Rect<2>(lowerBound, upperBound)); + DomainT<2> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<2>(distFusedIndexSpace)); + auto ADomain = runtime->get_index_space_domain(ctx, A_dense_run_0); + auto BDomain = runtime->get_index_space_domain(ctx, B_dense_run_0); + auto CDomain = runtime->get_index_space_domain(ctx, C_dense_run_0); + DomainPointColoring AColoring = DomainPointColoring(); + DomainPointColoring BColoring = DomainPointColoring(); + DomainPointColoring CColoring = DomainPointColoring(); + for (PointInDomainIterator<2> itr = PointInDomainIterator<2>(domain); itr.valid(); itr++) { + int64_t io = (*itr)[0]; + int64_t jo = (*itr)[1]; + Point<2> AStart = Point<2>((io * ((B1_dimension + (gx - 1)) / gx)), (jo * ((C2_dimension + (gy - 1)) / gy))); + Point<2> AEnd = Point<2>(TACO_MIN((io * ((B1_dimension + (gx - 1)) / gx) + ((B1_dimension + (gx - 1)) / gx - 1)), ADomain.hi()[0]), TACO_MIN((jo * ((C2_dimension + (gy - 1)) / gy) + ((C2_dimension + (gy - 1)) / gy - 1)), ADomain.hi()[1])); + Rect<2> ARect = Rect<2>(AStart, AEnd); + if (!ADomain.contains(ARect.lo) || !ADomain.contains(ARect.hi)) { + ARect = ARect.make_empty(); + } + AColoring[(*itr)] = ARect; + Point<1> BStart = Point<1>((io * ((B1_dimension + (gx - 1)) / gx))); + Point<1> BEnd = Point<1>(TACO_MIN((io * ((B1_dimension + (gx - 1)) / gx) + ((B1_dimension + (gx - 1)) / gx - 1)), BDomain.hi()[0])); + Rect<1> BRect = Rect<1>(BStart, BEnd); + if (!BDomain.contains(BRect.lo) || !BDomain.contains(BRect.hi)) { + BRect = BRect.make_empty(); + } + BColoring[(*itr)] = BRect; + Point<2> CStart = Point<2>(0, (jo * ((C2_dimension + (gy - 1)) / gy))); + Point<2> CEnd = Point<2>(TACO_MIN(B2_dimension, CDomain.hi()[0]), TACO_MIN((jo * ((C2_dimension + (gy - 1)) / gy) + ((C2_dimension + (gy - 1)) / gy - 1)), CDomain.hi()[1])); + Rect<2> CRect = Rect<2>(CStart, CEnd); + if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { + CRect = CRect.make_empty(); + } + CColoring[(*itr)] = CRect; + } + auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_DISJOINT_COMPLETE_KIND); + auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition, get_logical_region(A_vals)); + auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_ALIASED_COMPLETE_KIND); + LogicalPartition posPartB2 = copyPartition(ctx, runtime, B_dense_run_0_Partition, B2_pos); + LogicalPartition crdPartB2 = runtime->get_logical_partition(ctx, B2_crd, RectCompressedPosPartitionDownwards::apply(ctx, runtime, B2_crd.get_index_space(), posPartB2, B2_pos_parent, FID_RECT_1)); + auto B_vals_partition = copyPartition(ctx, runtime, crdPartB2, get_logical_region(B_vals)); + auto C_dense_run_0_Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_ALIASED_COMPLETE_KIND); + auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition, get_logical_region(C_vals)); + auto computePartitions = partitionPackForcomputeLegionConsMem(); + computePartitions.APartition.indicesPartitions = std::vector>(2); + computePartitions.APartition.denseLevelRunPartitions = std::vector(2); + computePartitions.APartition.valsPartition = A_vals_partition; + computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition; + computePartitions.BPartition.indicesPartitions = std::vector>(2); + computePartitions.BPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.BPartition.indicesPartitions[1].push_back(posPartB2); + computePartitions.BPartition.indicesPartitions[1].push_back(crdPartB2); + computePartitions.BPartition.valsPartition = B_vals_partition; + computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition; + computePartitions.CPartition.indicesPartitions = std::vector>(2); + computePartitions.CPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.CPartition.valsPartition = C_vals_partition; + computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition; + + return computePartitions; +} + +void task_2(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(); + PhysicalRegion B2_pos = regions[1]; + LogicalRegion B2_pos_parent = regions[1].get_logical_region(); + PhysicalRegion B2_crd = regions[2]; + LogicalRegion B2_crd_parent = regions[2].get_logical_region(); + PhysicalRegion B_vals = regions[3]; + LogicalRegion B_vals_parent = regions[3].get_logical_region(); + PhysicalRegion C_vals = regions[4]; + LogicalRegion C_vals_parent = regions[4].get_logical_region(); + + int64_t distFused = task->index_point[0]; + task_2Args* args = (task_2Args*)(task->args); + int64_t A2_dimension = args->A2_dimension; + int64_t B1_dimension = args->B1_dimension; + int64_t C2_dimension = args->C2_dimension; + int32_t gx = args->gx; + int32_t gy = args->gy; + + 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_crd_accessor = createAccessor(B2_crd, FID_COORD); + + int64_t io = getIndexPoint(task, 0); + int64_t jo = getIndexPoint(task, 1); + int64_t pointID1 = io; + int64_t pointID2 = pointID1 * gy + jo; + for (int64_t ii = 0; ii < ((B1_dimension + (gx - 1)) / gx); ii++) { + int64_t i = io * ((B1_dimension + (gx - 1)) / gx) + ii; + if (i >= B1_dimension) + continue; + + if (i >= (io + 1) * ((B1_dimension + (gx - 1)) / gx)) + continue; + + int64_t pointID3 = pointID2 * ((B1_dimension + (gx - 1)) / gx) + ii; + int64_t iA = i; + int64_t iB = i; + for (int64_t ji = 0; ji < ((C2_dimension + (gy - 1)) / gy); ji++) { + int64_t j = jo * ((C2_dimension + (gy - 1)) / gy) + ji; + if (j >= C2_dimension) + continue; + + if (j >= (jo + 1) * ((C2_dimension + (gy - 1)) / gy)) + continue; + + int64_t pointID4 = pointID3 * ((C2_dimension + (gy - 1)) / gy) + ji; + int64_t jA = iA * A2_dimension + j; + for (int64_t kB = B2_pos_accessor[Point<1>(i)].lo; kB < (B2_pos_accessor[Point<1>(i)].hi + 1); kB++) { + int64_t k = B2_crd_accessor[(kB * 1)]; + int64_t kC = k; + int64_t jC = kC * C2_dimension + j; + A_vals_rw_accessor[Point<2>(i, j)] = A_vals_rw_accessor[Point<2>(i, j)] + B_vals_ro_accessor[Point<1>(kB)] * C_vals_ro_accessor[Point<2>(k, j)]; + } + } + } +} + +void computeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionConsMem* partitionPack, int32_t gx, int32_t gy) { + int A2_dimension = A->dims[1]; + auto A_vals_parent = A->valsParent; + int B1_dimension = B->dims[0]; + RegionWrapper B2_crd = B->indices[1][1]; + auto B2_pos_parent = B->indicesParents[1][0]; + auto B2_crd_parent = B->indicesParents[1][1]; + auto B_vals_parent = B->valsParent; + int C2_dimension = C->dims[1]; + auto C_vals_parent = C->valsParent; + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + Point<2> lowerBound = Point<2>(0, 0); + Point<2> upperBound = Point<2>((gx - 1), (gy - 1)); + auto distFusedIndexSpace = runtime->create_index_space(ctx, Rect<2>(lowerBound, upperBound)); + DomainT<2> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<2>(distFusedIndexSpace)); + task_2Args taskArgsRaw2; + taskArgsRaw2.A2_dimension = A2_dimension; + taskArgsRaw2.B1_dimension = B1_dimension; + taskArgsRaw2.C2_dimension = C2_dimension; + taskArgsRaw2.gx = gx; + taskArgsRaw2.gy = gy; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); + IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(partitionPack->APartition.valsPartition, 0, READ_WRITE, EXCLUSIVE, A_vals_parent).add_field(FID_VAL)); + 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(partitionPack->CPartition.valsPartition, 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); + launcher.tag |= TACOMapper::BACKPRESSURE_TASK; + runtime->execute_index_space(ctx, launcher); + } void registerTacoTasks() { { TaskVariantRegistrar registrar(taskID(1), "task_1"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_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.set_leaf(); + Runtime::preregister_task_variant(registrar, "task_2"); + } } diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index 3b805b815..5a6477301 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -11,6 +11,7 @@ using namespace Legion; #include "leaf_kernels.cuh" typedef FieldAccessor> AccessorROdouble1; typedef FieldAccessor> AccessorROdouble2; +typedef FieldAccessor> AccessorRWdouble2; typedef ReductionAccessor,false,2,coord_t,Realm::AffineAccessor> AccessorReduceNonExcldouble2; typedef FieldAccessor> AccessorROint32_t1; typedef FieldAccessor,1,coord_t,Realm::AffineAccessor,1,coord_t>> AccessorRORect_1_1; @@ -21,6 +22,13 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t B1_dimension; + int64_t C2_dimension; + int32_t gx; + int32_t gy; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; @@ -200,6 +208,227 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* 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); +} + +partitionPackForcomputeLegionConsMem partitionForcomputeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx, int32_t gy) { + RegionWrapper A_vals = A->vals; + IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; + int B1_dimension = B->dims[0]; + RegionWrapper B2_pos = B->indices[1][0]; + RegionWrapper B2_crd = B->indices[1][1]; + auto B2_pos_parent = B->indicesParents[1][0]; + auto B2_crd_parent = B->indicesParents[1][1]; + RegionWrapper B_vals = B->vals; + IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; + auto B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); + auto B2_crd_accessor = createAccessor(B2_crd, FID_COORD); + int C2_dimension = C->dims[1]; + RegionWrapper C_vals = C->vals; + IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; + + B2_crd = legionMalloc( + ctx, + runtime, + B2_crd, + B2_crd_parent, + FID_COORD, + READ_ONLY + ); + B2_crd_accessor = createAccessor(B2_crd, FID_COORD); + B2_pos = legionMalloc( + ctx, + runtime, + B2_pos, + B2_pos_parent, + FID_RECT_1, + READ_ONLY + ); + B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); + + + Point<2> lowerBound = Point<2>(0, 0); + Point<2> upperBound = Point<2>((gx - 1), (gy - 1)); + auto distFusedIndexSpace = runtime->create_index_space(ctx, Rect<2>(lowerBound, upperBound)); + DomainT<2> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<2>(distFusedIndexSpace)); + auto ADomain = runtime->get_index_space_domain(ctx, A_dense_run_0); + auto BDomain = runtime->get_index_space_domain(ctx, B_dense_run_0); + auto CDomain = runtime->get_index_space_domain(ctx, C_dense_run_0); + DomainPointColoring AColoring = DomainPointColoring(); + DomainPointColoring BColoring = DomainPointColoring(); + DomainPointColoring CColoring = DomainPointColoring(); + for (PointInDomainIterator<2> itr = PointInDomainIterator<2>(domain); itr.valid(); itr++) { + int64_t io = (*itr)[0]; + int64_t jo = (*itr)[1]; + Point<2> AStart = Point<2>((io * ((B1_dimension + (gx - 1)) / gx)), (jo * ((C2_dimension + (gy - 1)) / gy))); + Point<2> AEnd = Point<2>(TACO_MIN((io * ((B1_dimension + (gx - 1)) / gx) + ((B1_dimension + (gx - 1)) / gx - 1)),ADomain.hi()[0]), TACO_MIN((jo * ((C2_dimension + (gy - 1)) / gy) + ((C2_dimension + (gy - 1)) / gy - 1)),ADomain.hi()[1])); + Rect<2> ARect = Rect<2>(AStart, AEnd); + if (!ADomain.contains(ARect.lo) || !ADomain.contains(ARect.hi)) { + ARect = ARect.make_empty(); + } + AColoring[(*itr)] = ARect; + Point<1> BStart = Point<1>((io * ((B1_dimension + (gx - 1)) / gx))); + Point<1> BEnd = Point<1>(TACO_MIN((io * ((B1_dimension + (gx - 1)) / gx) + ((B1_dimension + (gx - 1)) / gx - 1)),BDomain.hi()[0])); + Rect<1> BRect = Rect<1>(BStart, BEnd); + if (!BDomain.contains(BRect.lo) || !BDomain.contains(BRect.hi)) { + BRect = BRect.make_empty(); + } + BColoring[(*itr)] = BRect; + Point<2> CStart = Point<2>(0, (jo * ((C2_dimension + (gy - 1)) / gy))); + Point<2> CEnd = Point<2>(CDomain.hi()[0], TACO_MIN((jo * ((C2_dimension + (gy - 1)) / gy) + ((C2_dimension + (gy - 1)) / gy - 1)),CDomain.hi()[1])); + Rect<2> CRect = Rect<2>(CStart, CEnd); + if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { + CRect = CRect.make_empty(); + } + CColoring[(*itr)] = CRect; + } + auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_DISJOINT_COMPLETE_KIND); + auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition, get_logical_region(A_vals)); + auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_ALIASED_COMPLETE_KIND); + LogicalPartition posPartB2 = copyPartition(ctx, runtime, B_dense_run_0_Partition, B2_pos); + LogicalPartition crdPartB2 = runtime->get_logical_partition(ctx, B2_crd, RectCompressedPosPartitionDownwards::apply( + ctx, + runtime, + B2_crd.get_index_space(), + posPartB2, + B2_pos_parent, + FID_RECT_1 + )); + auto B_vals_partition = copyPartition(ctx, runtime, crdPartB2, get_logical_region(B_vals)); + auto C_dense_run_0_Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_ALIASED_COMPLETE_KIND); + auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition, get_logical_region(C_vals)); + auto computePartitions = partitionPackForcomputeLegionConsMem(); + computePartitions.APartition.indicesPartitions = std::vector>(2); + computePartitions.APartition.denseLevelRunPartitions = std::vector(2); + computePartitions.APartition.valsPartition = A_vals_partition; + computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition; + computePartitions.BPartition.indicesPartitions = std::vector>(2); + computePartitions.BPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.BPartition.indicesPartitions[1].push_back(posPartB2); + computePartitions.BPartition.indicesPartitions[1].push_back(crdPartB2); + computePartitions.BPartition.valsPartition = B_vals_partition; + computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition; + computePartitions.CPartition.indicesPartitions = std::vector>(2); + computePartitions.CPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.CPartition.valsPartition = C_vals_partition; + computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition; + + runtime->unmap_region(ctx, B2_crd); + runtime->unmap_region(ctx, B2_pos); + + return computePartitions; +} + +__global__ +void task_2DeviceKernel0(int32_t gy, int64_t io, int64_t jo, int64_t pointID2, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorRWdouble2 A_vals_rw_accessor, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t B1_dimension, int64_t C2_dimension, int32_t gx, int64_t distFused) { + + int64_t block = blockIdx.x; + int64_t thread = (threadIdx.x % (32)); + int64_t warp = (threadIdx.x / 32); + if (threadIdx.x >= 256) { + return; + } + + int64_t pointID3 = pointID2 * (((B1_dimension + (gx - 1)) / gx + 255) / 256) + block; + int64_t pointID4 = pointID3 * 8 + warp; + for (int64_t warp_row = 0; warp_row < 32; warp_row++) { + int64_t block_row = warp_row * 8 + warp; + int64_t ii = block * 256 + block_row; + int64_t i = io * ((B1_dimension + (gx - 1)) / gx) + ii; + if (i >= B1_dimension) + break; + + if (i >= (io + 1) * ((B1_dimension + (gx - 1)) / gx)) + break; + + int64_t pointID5 = pointID4 * 32 + warp_row; + int64_t pointID6 = pointID5 * 32 + thread; + for (int64_t thread_nz = 0; thread_nz < ((((B2_pos_accessor[Point<1>(i)].hi + 1) - B2_pos_accessor[Point<1>(i)].lo) + 31) / 32); thread_nz++) { + int64_t kposB = (thread_nz * 32 + thread) + B2_pos_accessor[Point<1>(i)].lo; + if (kposB < B2_pos_accessor[Point<1>(i)].lo || kposB >= B2_pos_accessor[Point<1>(i)].hi + 1) + break; + + int64_t k = B2_crd_accessor[kposB]; + int64_t pointID7 = pointID6 * ((((B2_pos_accessor[Point<1>(i)].hi + 1) - B2_pos_accessor[Point<1>(i)].lo) + 31) / 32) + thread_nz; + for (int64_t ji = 0; ji < ((C2_dimension + (gy - 1)) / gy); ji++) { + int64_t j = jo * ((C2_dimension + (gy - 1)) / gy) + ji; + if (j >= C2_dimension) + break; + + if (j >= (jo + 1) * ((C2_dimension + (gy - 1)) / gy)) + break; + + atomicAddWarp(A_vals_rw_accessor.ptr(Point<2>(i, j)), flattenPoint(A_vals_rw_accessor, Point<2>(i, j)), (B_vals_ro_accessor[Point<1>(kposB)] * C_vals_ro_accessor[Point<2>(k, j)])); + } + } + } +} + +void task_2(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(); + PhysicalRegion B2_pos = regions[1]; + LogicalRegion B2_pos_parent = regions[1].get_logical_region(); + PhysicalRegion B2_crd = regions[2]; + LogicalRegion B2_crd_parent = regions[2].get_logical_region(); + PhysicalRegion B_vals = regions[3]; + LogicalRegion B_vals_parent = regions[3].get_logical_region(); + PhysicalRegion C_vals = regions[4]; + LogicalRegion C_vals_parent = regions[4].get_logical_region(); + + int64_t distFused = task->index_point[0]; + task_2Args* args = (task_2Args*)(task->args); + int64_t B1_dimension = args->B1_dimension; + int64_t C2_dimension = args->C2_dimension; + int32_t gx = args->gx; + int32_t gy = args->gy; + + 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_crd_accessor = createAccessor(B2_crd, FID_COORD); + + if (runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).empty()) + return ; + + int64_t io = getIndexPoint(task, 0); + int64_t jo = getIndexPoint(task, 1); + int64_t pointID2 = io * gy + jo; + if ((((B1_dimension + (gx - 1)) / gx + 255) / 256) > 0) { + task_2DeviceKernel0<<<(((B1_dimension + (gx - 1)) / gx + 255) / 256), (32 * 8)>>>(gy, io, jo, pointID2, B2_pos_accessor, B2_crd_accessor, A_vals_rw_accessor, B_vals_ro_accessor, C_vals_ro_accessor, B1_dimension, C2_dimension, gx, distFused); + } +} + +void computeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionConsMem* partitionPack, int32_t gx, int32_t gy) { + auto A_vals_parent = A->valsParent; + int B1_dimension = B->dims[0]; + auto B2_pos_parent = B->indicesParents[1][0]; + auto B2_crd_parent = B->indicesParents[1][1]; + auto B_vals_parent = B->valsParent; + int C2_dimension = C->dims[1]; + auto C_vals_parent = C->valsParent; + + + Point<2> lowerBound = Point<2>(0, 0); + Point<2> upperBound = Point<2>((gx - 1), (gy - 1)); + auto distFusedIndexSpace = runtime->create_index_space(ctx, Rect<2>(lowerBound, upperBound)); + DomainT<2> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<2>(distFusedIndexSpace)); + task_2Args taskArgsRaw2; + taskArgsRaw2.B1_dimension = B1_dimension; + taskArgsRaw2.C2_dimension = C2_dimension; + taskArgsRaw2.gx = gx; + taskArgsRaw2.gy = gy; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); + IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(partitionPack->APartition.valsPartition, 0, READ_WRITE, EXCLUSIVE, A_vals_parent).add_field(FID_VAL)); + 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(partitionPack->CPartition.valsPartition, 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); + launcher.tag |= TACOMapper::UNTRACK_VALID_REGIONS; + launcher.tag |= TACOMapper::BACKPRESSURE_TASK; + runtime->execute_index_space(ctx, launcher); + } void registerTacoTasks() { { @@ -208,4 +437,10 @@ void registerTacoTasks() { registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_1"); } + { + TaskVariantRegistrar registrar(taskID(2), "task_2"); + registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC)); + registrar.set_leaf(); + Runtime::preregister_task_variant(registrar, "task_2"); + } } diff --git a/legion/spmm/taco-generated.cuh b/legion/spmm/taco-generated.cuh index 9e12eaef5..b4dbce896 100644 --- a/legion/spmm/taco-generated.cuh +++ b/legion/spmm/taco-generated.cuh @@ -8,10 +8,21 @@ struct partitionPackForcomputeLegion { LegionTensorPartition BPartition; }; +struct partitionPackForcomputeLegionConsMem { + LegionTensorPartition APartition; + LegionTensorPartition BPartition; + LegionTensorPartition CPartition; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegion* partitionPack, int32_t gx); + +partitionPackForcomputeLegionConsMem partitionForcomputeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx, int32_t gy); + + +void computeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionConsMem* partitionPack, int32_t gx, int32_t gy); void registerTacoTasks(); #endif // TACO_GENERATED_CUH diff --git a/legion/spmm/taco-generated.h b/legion/spmm/taco-generated.h index 6b952b5f0..df0d0600b 100644 --- a/legion/spmm/taco-generated.h +++ b/legion/spmm/taco-generated.h @@ -8,10 +8,21 @@ struct partitionPackForcomputeLegion { LegionTensorPartition BPartition; }; +struct partitionPackForcomputeLegionConsMem { + LegionTensorPartition APartition; + LegionTensorPartition BPartition; + LegionTensorPartition CPartition; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegion* partitionPack, int32_t gx); + +partitionPackForcomputeLegionConsMem partitionForcomputeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx, int32_t gy); + + +void computeLegionConsMem(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionConsMem* partitionPack, int32_t gx, int32_t gy); void registerTacoTasks(); #endif // TACO_GENERATED_H diff --git a/legion/spmttkrp/taco-generated.cpp b/legion/spmttkrp/taco-generated.cpp index 4a4912967..35486b1a5 100644 --- a/legion/spmttkrp/taco-generated.cpp +++ b/legion/spmttkrp/taco-generated.cpp @@ -430,13 +430,13 @@ void computeLegionDDS(Legion::Context ctx, Legion::Runtime* runtime, LegionTenso void registerTacoTasks() { { TaskVariantRegistrar registrar(taskID(1), "task_1"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_1"); } { TaskVariantRegistrar registrar(taskID(2), "task_2"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_2"); } diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index e157646a2..69890a75c 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -151,7 +151,7 @@ void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, if ((task.tag & UNTRACK_VALID_REGIONS) != 0 && this->untrackValidRegions) { for (size_t i = 0; i < task.regions.size(); i++) { auto &rg = task.regions[i]; - if (rg.privilege == READ_ONLY) { + if (rg.privilege == READ_ONLY && i == 4) { output.untracked_valid_regions.insert(i); } } @@ -579,15 +579,27 @@ void TACOMapper::slice_task(const Legion::Mapping::MapperContext ctx, // messes up the placement that we are going for with the index launches. This // implementation mirrors the standard slicing strategy of the default mapper. auto targets = this->select_targets_for_task(ctx, task); + bool recurse = ((task.tag & BACKPRESSURE_TASK) != 0) && (input.domain.get_volume() > 1); + if (input.domain.get_volume() == 1 && task.target_proc.exists()) { + targets = {task.target_proc}; + } + // We can't recurse here???? But we have to make sure that all of the blocks have size 1 I think. + recurse = false; + // bool recurse = false; + // std::cout << "RECURSE FOR TASK: " << task.get_task_name() << " " << task.index_domain << " " << recurse << " " << input.domain << std::endl; switch (input.domain.get_dim()) { #define BLOCK(DIM) \ case DIM: \ { \ Legion::DomainT point_space = input.domain; \ + auto targetSize = targets.size(); \ + if (((task.tag & BACKPRESSURE_TASK) != 0) && (input.domain.get_volume() > 1)) { \ + targetSize = input.domain.get_volume(); \ + } \ Legion::Point num_blocks = \ - default_select_num_blocks(targets.size(), point_space.bounds); \ + default_select_num_blocks(targetSize, point_space.bounds); \ this->default_decompose_points(point_space, targets, \ - num_blocks, false/*recurse*/, \ + num_blocks, recurse, \ stealing_enabled, output.slices); \ break; \ } @@ -597,6 +609,16 @@ void TACOMapper::slice_task(const Legion::Mapping::MapperContext ctx, taco_iassert(false); } } + std::cout << "SLICE_TASK for " + << Utilities::to_string(runtime, ctx, task, false /*include_index_point*/) + << " <" << task.get_unique_id() << ">" << std::endl; + std::cout << " INPUT: " << Utilities::to_string(runtime, ctx, input.domain) << std::endl; + std::cout << " OUTPUT:" << std::endl; + for (std::vector::const_iterator it = output.slices.begin(); + it != output.slices.end(); ++it) { + std::cout << " " << Utilities::to_string(runtime, ctx, it->domain) + << " -> " << it->proc << std::endl; + } } void TACOMapper::report_profiling(const MapperContext ctx, @@ -611,13 +633,16 @@ void TACOMapper::report_profiling(const MapperContext ctx, taco_iassert(prof->result == Realm::ProfilingMeasurements::OperationStatus::COMPLETED_SUCCESSFULLY); // Clean up after ourselves. delete prof; + + // std::cout << "Report profiling for task: " << task.get_task_name() << " " << task.get_unique_id() << " " << task.index_domain << " " << task.index_point << std::endl; + // Backpressured tasks are launched in a loop, and are kept on the originating processor. // So, we'll use orig_proc to index into the queue. auto& inflight = this->backPressureQueue[task.orig_proc]; MapperEvent event; // Find this task in the queue. for (auto it = inflight.begin(); it != inflight.end(); it++) { - if (it->id == task.get_unique_id()) { + if (it->id == std::make_pair(task.get_slice_domain(), task.get_context_index())) { event = it->event; inflight.erase(it); break; @@ -663,7 +688,9 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, { auto task = *it; bool schedule = true; - if ((task->tag & BACKPRESSURE_TASK) != 0) { + // TODO (rohany): Handle a case here when the total index space volume is also equal 1. + if (((task->tag & BACKPRESSURE_TASK) != 0) && !(task->is_index_space && task->get_slice_domain().get_volume() > 1)) { + // std::cout << "Visiting task: " << task->get_task_name() << " " << task->get_unique_id() << " " << task->index_domain << " " << task->index_point << " " << task->get_slice_domain() << std::endl; // See how many tasks we have in flight. Again, we use the orig_proc here // rather than target_proc to match with our heuristics for where serial task // launch loops go. @@ -683,7 +710,7 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, // Otherwise, we can schedule the task. Create a new event // and queue it up on the processor. this->backPressureQueue[task->orig_proc].push_back({ - .id = task->get_unique_id(), + .id = std::make_pair(task->get_slice_domain(), task->get_context_index()), .event = this->runtime->create_mapper_event(ctx), .schedTime = schedTime, }); @@ -692,6 +719,8 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, // Schedule tasks that are valid and have the target depth. if (schedule && (*it)->get_depth() == max_depth) { + auto dom = task->get_slice_domain(); + // std::cout << "Putting task in output map queue: " << Utilities::to_string(runtime, ctx, *task) << " " << task->index_domain << " " << (dom.exists() ? dom : Domain(Rect<1>(0, 0))) << std::endl; output.map_tasks.insert(*it); count++; } diff --git a/scripts/spbenchmark.py b/scripts/spbenchmark.py index 6ae684f06..04376e934 100755 --- a/scripts/spbenchmark.py +++ b/scripts/spbenchmark.py @@ -70,19 +70,58 @@ def getCommand(self, tensor, benchKind, procs): else: gpus = 4 # Allocate as much GPU memory as physically possible. - fbSize = "15.3G" + # TODO (rohany): Comment... + fbSize = "15G" + if procs >= 8: + # For some reason we get OOM errors at 32 GPUs if we try to take up too much of the + # GPU memory and we aren't able to allocate enough for GASNet-Ex. + fbSize = "14.5G" if benchKind == BenchmarkKind.SpMV: # We'll allocate slightly less memory in this case for CuSparse, as the cusparseHandle_t # takes up a non-trivial amount of memory. - fbSize = "15G" + fbSize = "14.5G" legionArgs += [ "-ll:gpu", str(gpus), "-ll:fsize", fbSize, "-pieces", str(procs), ] + + spmmArgs = [] + if self.gpu and benchKind == BenchmarkKind.SpMM: + spmmArgs += ["-consMem", "-tm:enable_backpressure", "-tm:backpressure_max_in_flight", "1", "-tm:untrack_valid_regions"] + if tensor.name in ["arabic-2005", "mycielskian19", "nlpkkt240"]: + # These inputs don't need any backpressuring, so overwrite the backpressuring args. + spmmArgs = ["-consMem", "-gx", str(procs), "-gy", "1"] + elif tensor.name == "it-2004": + if procs <= 8: + spmmArgs += ["-gx", "2", "-gy", "4"] + else: + spmmArgs += ["-gx", str(procs // 4), "-gy", "4"] + elif tensor.name == "sk-2005": + if procs <= 16: + spmmArgs += ["-gx", "4", "-gy", "4"] + else: + spmmArgs += ["-gx", str(procs // 4), "-gy", "4"] + elif tensor.name == "twitter7": + if procs <= 32: + spmmArgs += ["-gx", "4", "-gy", "8"] + else: + spmmArgs += ["-gx", str(procs // 8), "-gy", "4"] + elif tensor.name == "uk-2005": + if procs <= 8: + spmmArgs += ["-gx", "4", "-gy", "2"] + else: + spmmArgs += ["-gx", str(procs // 2), "-gy", "2"] + elif tensor.name == "webbase-2001": + # TODO (rohany): We can get 4,32 working on 1 node. + spmmArgs += ["-gx", "4", "-gy", "16"] + elif "kmer" in tensor.name: + spmmArgs += ["-gx", "8", "-gy", "8"] + else: + assert(False) assert(benchKind in args) - return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + spmmArgs def getBinary(self, benchKind): tacoDir = os.environ.get("TACO_BUILD_DIR", None) @@ -129,7 +168,7 @@ def getCommand(self, tensor, benchKind, nodes): BenchmarkKind.SpMTTKRP: ["-bench", "spmttkrp"], BenchmarkKind.SpInnerProd: ["-bench", "spinnerprod", "-tensorC", self.getShiftedTensor(tensor, 0)], } - lassenPrefix = ["jsrun", "-b", "rs", "-c", "1", "-r", "40", "-n", str(40 * nodes)] + lassenPrefix = ["jsrun", "-b", "rs", "-c", "1", "-r", "40", "-n", str(40 * nodes), "--core=none", "--core_delay=0"] commonArgs = ["-tensor", self.getCTFTensor(tensor), "-dims", ",".join([str(d) for d in tensor.dims]), "-n", str(self.niter), "-warmup", str(self.warmup)] ctfDir = os.environ.get("CTF_DIR", None) if (ctfDir is None): diff --git a/src/index_notation/provenance_graph.cpp b/src/index_notation/provenance_graph.cpp index 63f4679f8..c30cfd658 100644 --- a/src/index_notation/provenance_graph.cpp +++ b/src/index_notation/provenance_graph.cpp @@ -762,12 +762,7 @@ ir::Expr PosRelNode::recoverVariable(taco::IndexVar indexVar, ModeFunction segment_bounds = accessIterator.posBounds(parentPos); // positions should be with respect to entire array not just segment so don't need to offset variable when projecting. - ir::Expr project_result = ir::Load::make(coord_array, variableNames.at(getPosVar())); - - // but need to subtract parentvars start corodbound - ir::Expr parent_value = ir::Sub::make(project_result, parentCoordBounds[getParentVar()][0]); - - return parent_value; + return ir::Load::make(coord_array, variableNames.at(getPosVar())); } ir::Stmt PosRelNode::recoverChild(IndexVar indexVar, std::vector definedVarOrder, diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp index e2e173d07..2d7a3c1d1 100644 --- a/test/tests-distributed.cpp +++ b/test/tests-distributed.cpp @@ -1843,10 +1843,49 @@ TEST(distributed, legionSpMM) { .parallelize(thread, ParallelUnit::GPUThread, OutputRaceStrategy::Atomics) .communicate({A(i, j), B(i, k), C(k, j)}, fposo) ; + + auto memConsCPU = stmt.reorder({i, j, k}) + .distribute({i, j}, {io, jo}, {ii, ji}, Grid(gx, gy)) + .communicate({A(i, j), B(i, k), C(k, j)}, jo) + ; + + const int BATCHED_NNZ_PER_THREAD = 8; + const int BATCHED_NNZ_PER_WARP = BATCHED_NNZ_PER_THREAD * WARP_SIZE; + const int BATCHED_NNZ_PER_TB = BATCHED_NNZ_PER_THREAD * BLOCK_SIZE; + const int ROWS_PER_WARP = 1; + const int ROWS_PER_TB = ROWS_PER_WARP * BLOCK_SIZE; + + IndexVar fpos1("fpos1"), fpos2("fpos2"), thread_nz("thread_nz"), block_row("block_row"), warp_row("warp_row"); + auto memConsGPU = stmt.reorder({i, j, k}) + .distribute({i, j}, {io, jo}, {ii, ji}, Grid(gx, gy), taco::ParallelUnit::DistributedGPU) + .reorder({ii, k, ji}) + .split(ii, block, block_row, ROWS_PER_TB) + .split(block_row, warp_row, warp, BLOCK_SIZE / WARP_SIZE) + .pos(k, kpos, B(i, k)) + .split(kpos, thread_nz, thread, WARP_SIZE) + .reorder({block, warp, warp_row, thread, thread_nz}) + .parallelize(block, ParallelUnit::GPUBlock, OutputRaceStrategy::IgnoreRaces) + .parallelize(warp, ParallelUnit::GPUWarp, OutputRaceStrategy::IgnoreRaces) + // TODO (rohany): Is it possible to use a temporary here? + .parallelize(thread, ParallelUnit::GPUThread, OutputRaceStrategy::Atomics) + // .fuse(ii, k, f) + // .pos(f, fpos, B(i, k)) + // // TODO (rohany): This can probably be optimized further. + // .split(fpos, block, fpos1, BATCHED_NNZ_PER_TB) + // .split(fpos1, warp, fpos2, BATCHED_NNZ_PER_WARP) + // .split(fpos2, thread, thread_nz, BATCHED_NNZ_PER_THREAD) + // .reorder({block, warp, thread, thread_nz}) + // .parallelize(block, ParallelUnit::GPUBlock, OutputRaceStrategy::IgnoreRaces) + // .parallelize(warp, ParallelUnit::GPUWarp, OutputRaceStrategy::IgnoreRaces) + // .parallelize(thread, ParallelUnit::GPUThread, OutputRaceStrategy::Atomics) + .communicate({A(i, j), B(i, k), C(k, j)}, jo) + ; auto loweredCPU = lowerLegionSeparatePartitionCompute(cpuStmt, "computeLegion", false /* waitOnFutureMap */); auto loweredGPU = lowerLegionSeparatePartitionCompute(gpuStmt, "computeLegion", false /* waitOnFutureMap */); - ir::CodegenLegionC::compileToDirectory("../legion/spmm/", loweredCPU); - ir::CodegenLegionCuda::compileToDirectory("../legion/spmm/", loweredGPU); + auto loweredMemConsCPU = lowerLegionSeparatePartitionCompute(memConsCPU, "computeLegionConsMem", false /* waitOnFutureMap */); + auto loweredMemConsGPU = lowerLegionSeparatePartitionCompute(memConsGPU, "computeLegionConsMem", false /* waitOnFutureMap */); + ir::CodegenLegionC::compileToDirectory("../legion/spmm/", ir::Block::make(loweredCPU, loweredMemConsCPU)); + ir::CodegenLegionCuda::compileToDirectory("../legion/spmm/", ir::Block::make(loweredGPU, loweredMemConsGPU)); } TEST(distributed, legionSpSDDMM) {