From 883420603ff0cf3f00870b0f32d2e3783de86770 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Thu, 10 Feb 2022 15:29:45 -0800 Subject: [PATCH 01/12] *: support a batched schedule for SpMM that uses less memory --- include/taco/lower/lowerer_impl.h | 2 +- legion/include/taco_legion_header.h | 1 + legion/spmm/main.cpp | 23 +- legion/spmm/taco-generated.cpp | 197 +++++++++++++++- legion/spmm/taco-generated.cu | 262 +++++++++++++++++++++- legion/spmm/taco-generated.cuh | 11 + legion/spmm/taco-generated.h | 10 + src/lower/lowerer_impl.cpp | 119 +++++++--- src/lower/mode_format_rect_compressed.cpp | 3 +- src/type.cpp | 4 +- test/tests-distributed.cpp | 38 +++- 11 files changed, 625 insertions(+), 45 deletions(-) diff --git a/include/taco/lower/lowerer_impl.h b/include/taco/lower/lowerer_impl.h index 13cd07c1b..271a8f1b1 100644 --- a/include/taco/lower/lowerer_impl.h +++ b/include/taco/lower/lowerer_impl.h @@ -781,7 +781,7 @@ class LowererImpl : public util::Uncopyable { // construct an AffineProjection between the bottom-up partition of the top dense // level pack of `from` to the fully-dense level pack of `to`. It returns an empty // Expr (i.e. !ret.defined()) when `to` is not partitioned by `from`. - ir::Expr constructAffineProjection(Access& from, Access& to); + ir::Expr constructAffineProjection(Access& from, const Access& to); // construct a SparseGatherProjection between the partition of a // sparse level of `from` to the dense level pack of `to`. It returns diff --git a/legion/include/taco_legion_header.h b/legion/include/taco_legion_header.h index f3d290419..53f376ca8 100644 --- a/legion/include/taco_legion_header.h +++ b/legion/include/taco_legion_header.h @@ -16,6 +16,7 @@ enum TensorFields { }; const int TACO_TASK_BASE_ID = 10000; const int TACO_SHARD_BASE_ID = 1000; +const int TACO_PARTITION_COLOR_OFFSET = 10000; // A helper widget to treat LogicalRegions and PhysicalRegions the same. struct RegionWrapper { diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 75a954004..79ba77ed9 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -13,7 +13,7 @@ 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, batched = 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; @@ -24,6 +24,7 @@ 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); + parser.add_option_bool("-batched", batched); 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"; @@ -41,14 +42,28 @@ 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); + auto piecesIspace = runtime->create_index_space(ctx, Rect<1>(0, pieces - 1)); + auto piecesDom = runtime->get_index_space_domain(ctx, piecesIspace); - auto commPart = createSparseAliasingPartitions(ctx, runtime, A.vals.get_index_space(), pack.APartition.valsPartition.get_index_partition()); + partitionPackForcomputeLegion pack; + partitionPackForcomputeLegionBatched packBatched; + IndexPartition commPart; + if (batched) { + packBatched = partitionForcomputeLegionBatched(ctx, runtime, &A, &B, &C, pieces); + commPart = createSparseAliasingPartitions(ctx, runtime, A.vals.get_index_space(), packBatched.APartition.valsPartition.get_index_partition()); + } else { + pack = partitionForcomputeLegion(ctx, runtime, &A, &B, &C, pieces); + 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); 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); + if (batched) { + computeLegionBatched(ctx, runtime, &A, &B, &C, &packBatched, pieces); + } else { + 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 diff --git a/legion/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index 996552872..cd2796d64 100644 --- a/legion/spmm/taco-generated.cpp +++ b/legion/spmm/taco-generated.cpp @@ -18,6 +18,15 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t A2_dimension; + int64_t B1_dimension; + int64_t C2_dimension; + int32_t gx; + int64_t jo; + int64_t pointID1; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; @@ -31,6 +40,8 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg int C2_dimension = C->dims[1]; IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; + auto computePartitions = partitionPackForcomputeLegion(); + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; Point<1> lowerBound = Point<1>(0); @@ -66,7 +77,6 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg 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 computePartitions = partitionPackForcomputeLegion(); computePartitions.APartition.indicesPartitions = std::vector>(2); computePartitions.APartition.denseLevelRunPartitions = std::vector(2); computePartitions.APartition.valsPartition = A_vals_partition; @@ -106,7 +116,7 @@ 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); - int64_t pointID1 = io; + int64_t pointID1 = io + TACO_PARTITION_COLOR_OFFSET; #pragma omp parallel for schedule(dynamic, 128) for (int64_t ii = 0; ii < ((B1_dimension + (gx - 1)) / gx); ii++) { int64_t i = io * ((B1_dimension + (gx - 1)) / gx) + ii; @@ -162,14 +172,195 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* 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.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; runtime->execute_index_space(ctx, launcher); } + +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { + 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]; + RegionWrapper B_vals = B->vals; + IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; + int C2_dimension = C->dims[1]; + IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; + + auto computePartitions = partitionPackForcomputeLegionBatched(); + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; + Point<1> lowerBound = Point<1>(0); + Point<1> upperBound = Point<1>((gx - 1)); + auto ioIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(ioIndexSpace)); + 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<1> itr = PointInDomainIterator<1>(domain); itr.valid(); itr++) { + int64_t io = (*itr)[0]; + Point<2> AStart = Point<2>((io * ((B1_dimension + (gx - 1)) / gx)), (jo * 4)); + 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 * 4 + 3), 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; + } + 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), pointID1); + auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_DISJOINT_COMPLETE_KIND); + LogicalPartition posPartB2 = copyPartition(ctx, runtime, B_dense_run_0_Partition, B2_pos, pointID1); + 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, + pointID1 + )); + auto B_vals_partition = copyPartition(ctx, runtime, crdPartB2, get_logical_region(B_vals), pointID1); + 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; + } + + 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 io = 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; + int64_t jo = args->jo; + int64_t pointID1 = args->pointID1; + + 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 pointID2 = pointID1 * gx + io; + #pragma omp parallel for schedule(dynamic, 128) + 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 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; + for (int64_t ji = 0; ji < 4; ji++) { + int64_t j = jo * 4 + ji; + if (j >= C2_dimension) + continue; + + int64_t pointID4 = pointID3 * 4 + ji; + int64_t jA = iA * A2_dimension + j; + 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 computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx) { + int A2_dimension = A->dims[1]; + RegionWrapper A_vals = A->vals; + auto A_vals_parent = A->valsParent; + 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; + auto B_vals_parent = B->valsParent; + int C2_dimension = C->dims[1]; + RegionWrapper C_vals = C->vals; + auto C_vals_parent = C->valsParent; + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; + Point<1> lowerBound = Point<1>(0); + Point<1> upperBound = Point<1>((gx - 1)); + auto ioIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(ioIndexSpace)); + task_2Args taskArgsRaw2; + taskArgsRaw2.A2_dimension = A2_dimension; + taskArgsRaw2.B1_dimension = B1_dimension; + taskArgsRaw2.C2_dimension = C2_dimension; + taskArgsRaw2.gx = gx; + taskArgsRaw2.jo = jo; + taskArgsRaw2.pointID1 = pointID1; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); + IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(A_vals), pointID1), 0, READ_WRITE, EXCLUSIVE, A_vals_parent).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_pos), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_crd), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 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.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; + 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..f5435613f 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -21,6 +21,14 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t B2Size; + int64_t C2_dimension; + int32_t gx; + int64_t jo; + int64_t pointID1; +}; + partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; @@ -30,6 +38,9 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg auto B2_pos_parent = B->indicesParents[1][0]; RegionWrapper B_vals = B->vals; IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; + IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; + + auto computePartitions = partitionPackForcomputeLegion(); int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; @@ -39,6 +50,9 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(fposoIndexSpace)); DomainT<1> B2_crd_domain = runtime->get_index_space_domain(ctx, B2_crd.get_index_space()); DomainPointColoring B2_crd_coloring = DomainPointColoring(); + 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); for (PointInDomainIterator<1> itr = PointInDomainIterator<1>(domain); itr.valid(); itr++) { int64_t fposo = (*itr)[0]; Point<1> B2CrdStart = Point<1>((fposo * ((B2Size + (gx - 1)) / gx))); @@ -57,7 +71,8 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg B2_pos, B2_pos_parent, FID_RECT_1, - runtime->get_index_partition_color_space_name(ctx, B2_crd_index_part) + runtime->get_index_partition_color_space_name(ctx, B2_crd_index_part), + LEGION_ALIASED_INCOMPLETE_KIND ); IndexPartition posIndexPartB2 = densifyPartition(ctx, runtime, get_index_space(B2_pos), posSparsePartB2); LogicalPartition posPartB2 = runtime->get_logical_partition(ctx, B2_pos, posIndexPartB2); @@ -65,7 +80,6 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg IndexPartition BDenseRun0Partition = copyPartition(ctx, runtime, posPartB2, B_dense_run_0); IndexPartition ADenseRun0Partition = AffineProjection(0).apply(ctx, runtime, BDenseRun0Partition, A_dense_run_0); auto A_vals_partition = copyPartition(ctx, runtime, ADenseRun0Partition, get_logical_region(A_vals)); - auto computePartitions = partitionPackForcomputeLegion(); computePartitions.APartition.indicesPartitions = std::vector>(2); computePartitions.APartition.denseLevelRunPartitions = std::vector(2); computePartitions.APartition.valsPartition = A_vals_partition; @@ -81,7 +95,7 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg } __global__ -void task_1DeviceKernel0(int64_t B2Size, int32_t gx, int64_t* i_blockStarts, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorReduceNonExcldouble2 A_vals_red_accessor_non_excl, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t C2_dimension, int64_t fposo) { +void task_1DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_blockStarts, int64_t pointID1, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorReduceNonExcldouble2 A_vals_red_accessor_non_excl, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t C2_dimension) { int64_t block = blockIdx.x; int64_t thread = (threadIdx.x % (32)); @@ -90,7 +104,7 @@ void task_1DeviceKernel0(int64_t B2Size, int32_t gx, int64_t* i_blockStarts, Acc return; } - int64_t pointID2 = fposo * (((B2Size + (gx - 1)) / gx + 511) / 512) + block; + int64_t pointID2 = pointID1 * (((B2Size + (gx - 1)) / gx + 511) / 512) + block; int64_t pointID3 = pointID2 * 8 + warp; int64_t pointID4 = pointID3 * 32 + thread; for (int64_t dense_b = 0; dense_b < ((C2_dimension + 31) / 32); dense_b++) { @@ -166,8 +180,9 @@ void task_1(const Task* task, const std::vector& regions, Contex (((B2Size + (gx - 1)) / gx + 511) / 512), B2CrdDomain.bounds.lo ); - if (((((B2Size + (gx - 1)) / gx + 511) / 512)) > 0) { - task_1DeviceKernel0<<<(((B2Size + (gx - 1)) / gx + 511) / 512), (32 * 8)>>>(B2Size, gx, i_blockStarts, B2_pos_accessor, B2_crd_accessor, A_vals_red_accessor_non_excl, B_vals_ro_accessor, C_vals_ro_accessor, C2_dimension, fposo); + int64_t pointID1 = fposo + TACO_PARTITION_COLOR_OFFSET; + if ((((B2Size + (gx - 1)) / gx + 511) / 512) > 0) { + task_1DeviceKernel0<<<(((B2Size + (gx - 1)) / gx + 511) / 512), (32 * 8)>>>(B2Size, fposo, gx, i_blockStarts, pointID1, B2_pos_accessor, B2_crd_accessor, A_vals_red_accessor_non_excl, B_vals_ro_accessor, C_vals_ro_accessor, C2_dimension); } } @@ -198,9 +213,238 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* 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.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; runtime->execute_index_space(ctx, launcher); } + +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { + 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]; + + auto computePartitions = partitionPackForcomputeLegionBatched(); + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; + Point<1> lowerBound = Point<1>(0); + Point<1> upperBound = Point<1>((gx - 1)); + auto fposoIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(fposoIndexSpace)); + DomainT<1> B2_crd_domain = runtime->get_index_space_domain(ctx, B2_crd.get_index_space()); + DomainPointColoring B2_crd_coloring = DomainPointColoring(); + 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 CColoring = DomainPointColoring(); + for (PointInDomainIterator<1> itr = PointInDomainIterator<1>(domain); itr.valid(); itr++) { + int64_t fposo = (*itr)[0]; + Point<1> B2CrdStart = Point<1>((fposo * ((B2Size + (gx - 1)) / gx))); + Point<1> B2CrdEnd = Point<1>(TACO_MIN((fposo * ((B2Size + (gx - 1)) / gx) + ((B2Size + (gx - 1)) / gx - 1)),B2_crd_domain.bounds.hi[0])); + Rect<1> B2CrdRect = Rect<1>(B2CrdStart, B2CrdEnd); + if (!B2_crd_domain.contains(B2CrdRect.lo) || !B2_crd_domain.contains(B2CrdRect.hi)) { + B2CrdRect = B2CrdRect.make_empty(); + } + B2_crd_coloring[(*itr)] = B2CrdRect; + Point<2> CStart = Point<2>(0, (jo * 4)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 4 + 3),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; + } + IndexPartition B2_crd_index_part = runtime->create_index_partition( + ctx, + B2_crd.get_index_space(), + domain, + B2_crd_coloring, + LEGION_COMPUTE_KIND, + pointID1 + ); + 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), + LEGION_ALIASED_INCOMPLETE_KIND + ); + IndexPartition posIndexPartB2 = densifyPartition(ctx, runtime, get_index_space(B2_pos), posSparsePartB2, pointID1); + LogicalPartition posPartB2 = runtime->get_logical_partition(ctx, B2_pos, posIndexPartB2); + LogicalPartition BValsLogicalPart = copyPartition(ctx, runtime, B2_crd_part, B_vals, pointID1); + IndexPartition BDenseRun0Partition = copyPartition(ctx, runtime, posPartB2, B_dense_run_0); + IndexPartition ADenseRun0Partition = AffineProjection(0).apply(ctx, runtime, BDenseRun0Partition, A_dense_run_0); + auto A_vals_partition = copyPartition(ctx, runtime, ADenseRun0Partition, get_logical_region(A_vals), pointID1); + IndexPartition CDenseRun0Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND); + auto C_vals_partition = copyPartition(ctx, runtime, CDenseRun0Partition, get_logical_region(C_vals), pointID1); + computePartitions.APartition.indicesPartitions = std::vector>(2); + computePartitions.APartition.denseLevelRunPartitions = std::vector(2); + computePartitions.APartition.valsPartition = A_vals_partition; + computePartitions.APartition.denseLevelRunPartitions[0] = ADenseRun0Partition; + 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(B2_crd_part); + computePartitions.BPartition.valsPartition = BValsLogicalPart; + computePartitions.BPartition.denseLevelRunPartitions[0] = BDenseRun0Partition; + computePartitions.CPartition.indicesPartitions = std::vector>(2); + computePartitions.CPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.CPartition.valsPartition = C_vals_partition; + computePartitions.CPartition.denseLevelRunPartitions[0] = CDenseRun0Partition; + } + + return computePartitions; +} + +__global__ +void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_blockStarts, int64_t pointID2, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorReduceNonExcldouble2 A_vals_red_accessor_non_excl, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t C2_dimension, int64_t jo, int64_t pointID1) { + + 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 * (((B2Size + (gx - 1)) / gx + 2047) / 2048) + block; + int64_t pointID4 = pointID3 * 8 + warp; + int64_t pointID5 = pointID4 * 32 + thread; + int64_t pB2_begin = i_blockStarts[block]; + int64_t pB2_end = i_blockStarts[(block + 1)]; + int64_t fposi2 = thread * 8; + int64_t fposi1 = warp * 256 + fposi2; + int64_t fposi = block * 2048 + fposi1; + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx) + fposi; + int64_t i_pos = taco_binarySearchBefore(B2_pos_accessor, pB2_begin, pB2_end, fposB); + int64_t i = i_pos; + for (int64_t thread_nz = 0; thread_nz < 8; thread_nz++) { + int64_t fposi2 = thread * 8 + thread_nz; + int64_t fposi1 = warp * 256 + fposi2; + int64_t fposi = block * 2048 + fposi1; + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx) + fposi; + if (fposB >= (fposo + 1) * ((B2Size + (gx - 1)) / gx)) + break; + + if (fposB >= B2Size) + break; + + int64_t f = B2_crd_accessor[fposB]; + while (!(B2_pos_accessor[i_pos].contains(fposB))) { + i_pos = i_pos + 1; + i = i_pos; + } + for (int64_t ji = 0; ji < 4; ji++) { + int64_t j = jo * 4 + ji; + if (j >= C2_dimension) + break; + + atomicAddWarp(A_vals_red_accessor_non_excl.ptr(Point<2>(i, j)), flattenPoint(A_vals_red_accessor_non_excl, Point<2>(i, j)), (B_vals_ro_accessor[Point<1>(fposB)] * C_vals_ro_accessor[Point<2>(f, 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 fposo = task->index_point[0]; + task_2Args* args = (task_2Args*)(task->args); + int64_t B2Size = args->B2Size; + int64_t C2_dimension = args->C2_dimension; + int32_t gx = args->gx; + int64_t jo = args->jo; + int64_t pointID1 = args->pointID1; + + auto B_vals_ro_accessor = createAccessor(B_vals, FID_VAL); + auto C_vals_ro_accessor = createAccessor(C_vals, FID_VAL); + auto A_vals_red_accessor_non_excl = createAccessor(A_vals, FID_VAL, LEGION_REDOP_SUM_FLOAT64); + 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 ; + + DomainT<1> B2PosDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_pos)); + DomainT<1> B2CrdDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)); + Legion::DeferredBuffer buf = Legion::DeferredBuffer(Rect<1>(0, (((B2Size + (gx - 1)) / gx + 2047) / 2048)), Legion::Memory::Kind::GPU_FB_MEM); + int64_t* i_blockStarts = buf.ptr(0); + taco_binarySearchBeforeBlockLaunch( + B2_pos_accessor, + i_blockStarts, + B2PosDomain.bounds.lo, + B2PosDomain.bounds.hi, + 2048, + 256, + (((B2Size + (gx - 1)) / gx + 2047) / 2048), + B2CrdDomain.bounds.lo + ); + int64_t pointID2 = pointID1 * gx + fposo; + if ((((B2Size + (gx - 1)) / gx + 2047) / 2048) > 0) { + task_2DeviceKernel0<<<(((B2Size + (gx - 1)) / gx + 2047) / 2048), (32 * 8)>>>(B2Size, fposo, gx, i_blockStarts, pointID2, B2_pos_accessor, B2_crd_accessor, A_vals_red_accessor_non_excl, B_vals_ro_accessor, C_vals_ro_accessor, C2_dimension, jo, pointID1); + } +} + +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx) { + RegionWrapper A_vals = A->vals; + auto A_vals_parent = A->valsParent; + 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; + auto B_vals_parent = B->valsParent; + int C2_dimension = C->dims[1]; + RegionWrapper C_vals = C->vals; + auto C_vals_parent = C->valsParent; + + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; + + for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; + Point<1> lowerBound = Point<1>(0); + Point<1> upperBound = Point<1>((gx - 1)); + auto fposoIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(fposoIndexSpace)); + task_2Args taskArgsRaw2; + taskArgsRaw2.B2Size = B2Size; + taskArgsRaw2.C2_dimension = C2_dimension; + taskArgsRaw2.gx = gx; + taskArgsRaw2.jo = jo; + taskArgsRaw2.pointID1 = pointID1; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); + IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(A_vals), pointID1), 0, LEGION_REDOP_SUM_FLOAT64, LEGION_SIMULTANEOUS, A_vals_parent).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_pos), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_crd), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); + launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; + runtime->execute_index_space(ctx, launcher); + + } +} void registerTacoTasks() { { TaskVariantRegistrar registrar(taskID(1), "task_1"); @@ -208,4 +452,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..adf2ba615 100644 --- a/legion/spmm/taco-generated.cuh +++ b/legion/spmm/taco-generated.cuh @@ -8,10 +8,21 @@ struct partitionPackForcomputeLegion { LegionTensorPartition BPartition; }; +struct partitionPackForcomputeLegionBatched { + 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); + +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); + + +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx); void registerTacoTasks(); #endif // TACO_GENERATED_CUH diff --git a/legion/spmm/taco-generated.h b/legion/spmm/taco-generated.h index 6b952b5f0..ad9a79f7b 100644 --- a/legion/spmm/taco-generated.h +++ b/legion/spmm/taco-generated.h @@ -8,10 +8,20 @@ struct partitionPackForcomputeLegion { LegionTensorPartition BPartition; }; +struct partitionPackForcomputeLegionBatched { + LegionTensorPartition APartition; + LegionTensorPartition BPartition; +}; + 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); + +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); + + +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx); void registerTacoTasks(); #endif // TACO_GENERATED_H diff --git a/src/lower/lowerer_impl.cpp b/src/lower/lowerer_impl.cpp index b9c8eccdd..b6a54a595 100644 --- a/src/lower/lowerer_impl.cpp +++ b/src/lower/lowerer_impl.cpp @@ -1929,8 +1929,8 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, Expr rhs; if (i == 0) { // In the first level of the iteration, the iteration space identifier is - // the variable itself. - rhs = this->indexVarToExprMap[this->definedIndexVarsExpanded[i]]; + // the variable itself, offset by an initial value. + rhs = ir::Add::make(this->indexVarToExprMap[this->definedIndexVarsExpanded[i]], ir::Symbol::make("TACO_PARTITION_COLOR_OFFSET")); } else { // Otherwise, we construct the variable from the prior level iteration space point. auto var = this->indexVarToExprMap[this->definedIndexVarsExpanded[i]]; @@ -2183,6 +2183,53 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, coloringLoopBody.push_back(ir::VarDecl::make(this->indexVarToExprMap[forall.getIndexVar()], ir::Load::make(ir::Deref::make(domainIter, pointT), 0))); coloringLoopBody.push_back(posIter.getCreatePosColoringEntry(ir::Deref::make(domainIter, Auto), lower, upper)); + + + std::map extraColorings; + std::vector colorings; + std::vector initColorings; + for (auto &t : forall.getTransfers()) { + auto c = ir::Var::make(t.getAccess().getTensorVar().getName() + "Coloring", DomainPointColoring); + initColorings.push_back( + ir::VarDecl::make(c, ir::Call::make(DomainPointColoring.getName(), {}, DomainPointColoring))); + colorings.push_back(c); + } + // TODO (rohany): Hacking... + for (auto& t : forall.getTransfers()) { + auto tv = t.getAccess().getTensorVar(); + auto domain = ir::Var::make(tv.getName() + "Domain", Auto); + auto runs = DenseFormatRuns(t.getAccess(), this->iterators); + auto ispace = ir::GetProperty::makeDenseLevelRun(this->tensorVars[tv], 0); + partitioningStmts.push_back(ir::VarDecl::make(domain, ir::Call::make("runtime->get_index_space_domain", {ctx, ispace}, Auto))); + domains[t.getAccess().getTensorVar()] = domain; + } + for (size_t idx = 0; idx < forall.getTransfers().size(); idx++) { + auto& t = forall.getTransfers()[idx]; + auto& tv = t.getAccess().getTensorVar(); + auto n = tv.getName(); + if (t.getAccess() == posAccess) { + continue; + } + auto affineProjection = this->constructAffineProjection(posAccess, t.getAccess()); + bool isPartitionedByOuterVars = false; + // std::set considering = this->varsInScope[this->curDistVar]; + std::set considering = this->definedIndexVars; + considering.erase(this->curDistVar); + for (auto ivar : t.getAccess().getIndexVars()) { + if (this->anyParentInSet(ivar, considering)) { + isPartitionedByOuterVars = true; + } + } + if (isPartitionedByOuterVars && !affineProjection.defined()) { + partitioningStmts.push_back(initColorings[idx]); + std::set fullyRepl; + coloringLoopBody.push_back(ir::Block::make(this->createDomainPointColorings(forall, domainIter, domains, fullyRepl, colorings, tv))); + extraColorings[tv] = colorings[idx]; + } + // TODO (rohany): As described below, we aren't going to be able to compute a tight bound + // on the subspaces here. + // taco_iassert(!(isPartitionedByOuterVars && affineProjection.defined())); + } auto coloringLoop = ir::For::make( domainIter, ir::Call::make(pointInDimT.getName(), {domain}, pointInDimT), @@ -2190,6 +2237,19 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, 1 /* increment -- hack to get ++ */, ir::Block::make(coloringLoopBody) ); + + // TODO (rohany): We need to add coloring and partition ops for tensors that are being partitioned + // by variables that are in scope but not present in the projection from the pos tensor. I.e. + // if we had A(i, j) = B(i, k) * C(k, j), ordering j -> i -> k and do a pos on B(i, k). In this + // case, the C tensor is not partitioned by the projection from B, but needs to be partitioned + // by the current value of j. + // TODO (rohany): It's unclear to me what we would do if we had a partition implied from a projection + // and by in scope variables. For example, consider the contraction A(i, j) = B(i, k) * C(i, j, k), + // with loop ordering j -> i -> k and a pos on B(i, k). In this case, the C tensor is partitioned by + // both the partition of B implied by the position split and the current value of j. This seems like + // a more general case of what we have now, where we'd want to form a coloring of the parent partitioned + // dimensions of C, and take a product of some sort against dimensions implied by the projection from B. + partitioningStmts.push_back(coloringLoop); partitioningStmts.push_back(posIter.getFinalizePosColoring()); @@ -2286,7 +2346,7 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, auto denseRunPartition = ir::Var::make(posTensor.getName() + "DenseRun0Partition", IndexPartition); DenseFormatRuns posTensorFormatRuns(posAccess, this->iterators); if (!posTensorFormatRuns.runs.empty() && util::contains(posTensorFormatRuns.runs[0].levels, 0)) { - partitioningStmts.push_back(ir::VarDecl::make(denseRunPartition, ir::Call::make("copyPartition", maybeAddPartColor({ctx, runtime, upwardsPart, denseRun}), IndexPartition))); + partitioningStmts.push_back(ir::VarDecl::make(denseRunPartition, ir::Call::make("copyPartition", {ctx, runtime, upwardsPart, denseRun}, IndexPartition))); tensorDenseRunPartitions[posTensor][0] = denseRunPartition; } @@ -2368,27 +2428,31 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, ir::Expr createDenseRunPart; auto firstDenseRun = ir::GetProperty::makeDenseLevelRun(tvExpr, 0); - // Attempt to create different projections from the pos tensor - // into the target tensor. - auto affineProjection = this->constructAffineProjection(posAccess, tvAccess); - auto sparseGatherProjection = this->constructSparseGatherProjection(posAccess, tvAccess, posMode); - if (affineProjection.defined()) { - createDenseRunPart = ir::MethodCall::make(affineProjection, "apply", maybeAddPartColor({ctx, runtime, denseRunPartition, firstDenseRun}), false /* deref */, IndexPartition); - } else if (sparseGatherProjection.defined() && !isFusedPos) { - // We can't apply a SparseGatherProjection if the position space loop is fused. - // TODO (rohany): This is still hardcoded for the RectCompressedModeFormat. - auto rcmf = posIter.getMode().getModeFormat().as(); - auto pack = posIter.getMode().getModePack(); - auto crdParent = rcmf->getRegion(pack, RectCompressedModeFormat::CRD_PARENT); - auto crdAcc = rcmf->getAccessor(pack, RectCompressedModeFormat::CRD, ir::RO).as(); - auto crdField = crdAcc->accessorArgs.field; - auto crdPart = tensorLogicalPartitions[posTensor][posLevel][1]; - createDenseRunPart = ir::MethodCall::make(sparseGatherProjection, "apply", maybeAddPartColor( - {ctx, runtime, crdParent, crdPart, crdField, firstDenseRun}), false /* deref */, - IndexPartition); + if (util::contains(extraColorings, tv)) { + createDenseRunPart = ir::Call::make("runtime->create_index_partition", {ctx, firstDenseRun, domain, extraColorings[tv], computePart}, Auto); } else { - // If we can't make any projections, then this tensor is fully replicated. - continue; + // Attempt to create different projections from the pos tensor + // into the target tensor. + auto affineProjection = this->constructAffineProjection(posAccess, tvAccess); + auto sparseGatherProjection = this->constructSparseGatherProjection(posAccess, tvAccess, posMode); + if (affineProjection.defined()) { + createDenseRunPart = ir::MethodCall::make(affineProjection, "apply", {ctx, runtime, denseRunPartition, firstDenseRun}, false /* deref */, IndexPartition); + } else if (sparseGatherProjection.defined() && !isFusedPos) { + // We can't apply a SparseGatherProjection if the position space loop is fused. + // TODO (rohany): This is still hardcoded for the RectCompressedModeFormat. + auto rcmf = posIter.getMode().getModeFormat().as(); + auto pack = posIter.getMode().getModePack(); + auto crdParent = rcmf->getRegion(pack, RectCompressedModeFormat::CRD_PARENT); + auto crdAcc = rcmf->getAccessor(pack, RectCompressedModeFormat::CRD, ir::RO).as(); + auto crdField = crdAcc->accessorArgs.field; + auto crdPart = tensorLogicalPartitions[posTensor][posLevel][1]; + createDenseRunPart = ir::MethodCall::make(sparseGatherProjection, "apply", maybeAddPartColor( + {ctx, runtime, crdParent, crdPart, crdField, firstDenseRun}), false /* deref */, + IndexPartition); + } else { + // If we can't make any projections, then this tensor is fully replicated. + continue; + } } // TODO (rohany): Deduplicate this out into a helper function used by the standard @@ -2636,7 +2700,7 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, // First, instantiate an instance of the struct in a pointer. auto structPack = ir::Var::make("computePartitions", Auto); - returnPartitionStatements.push_back(ir::VarDecl::make(structPack, ir::makeConstructor(structTy, {}))); + this->header.push_back(ir::VarDecl::make(structPack, ir::makeConstructor(structTy, {}))); for (auto& t : this->tensorVarOrdering) { auto tensor = this->tensorVars[t]; @@ -5709,6 +5773,9 @@ Expr LowererImpl::projectCanonicalSpaceToWindowedPosition(Iterator iterator, ir: } bool LowererImpl::anyParentInSet(IndexVar var, std::set& s) { + if (util::contains(s, var)) { + return true; + } auto children = this->provGraph.getChildren(var); for (auto c : children) { if (util::contains(s, c)) { @@ -6308,6 +6375,7 @@ std::vector LowererImpl::lowerIndexLaunch( ir::Expr tag; if (!util::contains(regionsAccessed, regions[i].region.as()->toHashable()) && !(this->isPlacementCode && size_t(this->distLoopDepth + 1) == this->placements.size())) { tag = virtualMap; + } else { taskReadsAnyVars |= true; } if (util::contains(tensorLogicalPartitions, tv)) { @@ -6666,6 +6734,7 @@ std::vector LowererImpl::lowerSerialTaskLoop( ir::Expr tag; if (!util::contains(regionsAccessed, regions[i].region.as()->toHashable()) && !(this->isPlacementCode && size_t(this->distLoopDepth + 1) == this->placements.size())) { tag = virtualMap; + } else { taskReadsAnyVars |= true; } @@ -6981,7 +7050,7 @@ std::vector LowererImpl::getAllNeededParentPositions(Iterator &iter) { return {rev.begin(), rev.end()}; } -ir::Expr LowererImpl::constructAffineProjection(Access &from, Access &to) { +ir::Expr LowererImpl::constructAffineProjection(Access &from, const Access &to) { auto fromDenseRuns = DenseFormatRuns(from, this->iterators); auto toDenseRuns = DenseFormatRuns(to, this->iterators); diff --git a/src/lower/mode_format_rect_compressed.cpp b/src/lower/mode_format_rect_compressed.cpp index 9bd6e9146..29f7f07e9 100644 --- a/src/lower/mode_format_rect_compressed.cpp +++ b/src/lower/mode_format_rect_compressed.cpp @@ -829,7 +829,8 @@ ModeFunction RectCompressedModeFormat::partitionPosFromCrd(Mode mode, ir::Expr c this->getRegion(pack, POS), this->getRegion(pack, POS_PARENT), fidRect1, - ir::Call::make("runtime->get_index_partition_color_space_name", {ir::ctx, crdIndexPartition}, Auto) + ir::Call::make("runtime->get_index_partition_color_space_name", {ir::ctx, crdIndexPartition}, Auto), + ir::Symbol::make("LEGION_ALIASED_INCOMPLETE_KIND"), }, Auto ); diff --git a/src/type.cpp b/src/type.cpp index a9234193d..32a274ce3 100644 --- a/src/type.cpp +++ b/src/type.cpp @@ -55,7 +55,9 @@ Datatype max_type(Datatype a, Datatype b) { if (a == b) { return a; } - else if (a.isComplex() || b.isComplex()) { + if (a.getKind() == Datatype::CppType || b.getKind() == Datatype::CppType) { + return Datatype::CppType; + } else if (a.isComplex() || b.isComplex()) { if (a == Complex128 || b == Complex128 || a == Float64 || b == Float64) { return Complex128; } diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp index 6be674702..97d625625 100644 --- a/test/tests-distributed.cpp +++ b/test/tests-distributed.cpp @@ -1817,8 +1817,8 @@ TEST(distributed, legionSpMM) { .communicate({A(i, j), B(i, k), C(k, j)}, io) ; IndexVar fpos("fpos"), fposo("fposo"), fposi("fposi"), block("block"), - warp("warp"), thread("thread"), fposi1("fposi1"), nnz("nnz"), - dense_ub("dense_ub"), dense_b("dense_b"); + warp("warp"), thread("thread"), fposi1("fposi1"), fposi2("fposi2"), nnz("nnz"), + dense_ub("dense_ub"), dense_b("dense_b"), thread_nz("thread_nz"); const int NNZ_PER_WARP = 64; const int BLOCK_SIZE = 256; const int WARP_SIZE = 32; @@ -1845,8 +1845,38 @@ TEST(distributed, legionSpMM) { ; 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); + + const int BATCH_SIZE = 4; + // Schedule to conserve memory (i.e. not replicate C onto all nodes). + auto cpuCons = stmt.split(j, jo, ji, BATCH_SIZE) + .reorder({jo, i, k, ji}) + .distribute({i}, {io}, {ii}, Grid(gx)) + .parallelize(ii, ParallelUnit::CPUThread, OutputRaceStrategy::NoRaces) + .communicate({A(i, j), B(i, k), C(k, j)}, io) + ; + + 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; + auto cpuConsLowered = lowerLegionSeparatePartitionCompute(cpuCons, "computeLegionBatched", false /* waitOnFutureMap */); + auto gpuCons = stmt.split(j, jo, ji, BATCH_SIZE) + .reorder({jo, i, k, ji}) + .fuse(i, k, f) + .pos(f, fpos, B(i, k)) + .distribute({fpos}, {fposo}, {fposi}, Grid(gx), taco::ParallelUnit::DistributedGPU) + // TODO (rohany): This can probably be optimized further. + .split(fposi, block, fposi1, BATCHED_NNZ_PER_TB) + .split(fposi1, warp, fposi2, BATCHED_NNZ_PER_WARP) + .split(fposi2, 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)}, fposo) + ; + auto gpuConsLowered = lowerLegionSeparatePartitionCompute(gpuCons, "computeLegionBatched", false /* waitOnFutureMap */); + ir::CodegenLegionC::compileToDirectory("../legion/spmm/", ir::Block::make(loweredCPU, cpuConsLowered)); + ir::CodegenLegionCuda::compileToDirectory("../legion/spmm/", ir::Block::make(loweredGPU, gpuConsLowered)); } TEST(distributed, legionSpSDDMM) { From 7e4601d57b6473f2872b5aa5960da03bc02d9d0b Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Fri, 11 Feb 2022 15:42:47 -0800 Subject: [PATCH 02/12] big fixes, support partitioning by affine and projection --- legion/include/taco_legion_header.h | 12 +- legion/include/taco_mapper.h | 7 + legion/spmm/main.cpp | 1 + legion/spmm/taco-generated.cpp | 251 ++++++++++++++++++-------- legion/spmm/taco-generated.cu | 10 +- legion/spmm/taco-generated.h | 2 + legion/src/taco_legion_header.cpp | 16 +- legion/src/taco_mapper.cpp | 266 +++++++++++++++++++++++++++- src/lower/lowerer_impl.cpp | 28 ++- test/tests-distributed.cpp | 11 +- 10 files changed, 517 insertions(+), 87 deletions(-) diff --git a/legion/include/taco_legion_header.h b/legion/include/taco_legion_header.h index 53f376ca8..5949eca1d 100644 --- a/legion/include/taco_legion_header.h +++ b/legion/include/taco_legion_header.h @@ -202,7 +202,7 @@ class AffineProjection { // Access the mapped index of this projection. int operator[](size_t i) const; // Apply the projection to an index partition. - Legion::IndexPartition apply(Legion::Context ctx, Legion::Runtime* runtime, Legion::IndexPartition part, Legion::IndexSpace ispace, Legion::Color color = LEGION_AUTO_GENERATE_ID); + Legion::IndexPartition apply(Legion::Context ctx, Legion::Runtime* runtime, Legion::IndexPartition part, Legion::IndexSpace ispace, Legion::DomainPointColoring coloring = Legion::DomainPointColoring(), Legion::Color color = LEGION_AUTO_GENERATE_ID); // Apply the projection to a domain point. outputBounds contains the corresponding // bounds for the output point (i.e. zero, n). It is used when positions in the // output are not specified by the AffineProjection, as in situations where a lower @@ -210,11 +210,21 @@ class AffineProjection { // are not specified by the projection, they take the value given by outputBounds at // that dimension. Legion::DomainPoint apply(Legion::DomainPoint point, Legion::DomainPoint outputBounds); + + AffineProjection addOverrides(std::vector overs) { + for (auto it : overs) { + this->overrides.push_back(it); + } + return *this; + } + template + AffineProjection addOverrides(const Overs&... overs) { return addOverrides(std::vector{overs...}); } // Value to represent the \bot value of projection, i.e. an index // that does not map to other indices. const static int BOT; private: std::vector projs; + std::vector overrides; }; // Helper class to perform partitions for the downwards partition path for diff --git a/legion/include/taco_mapper.h b/legion/include/taco_mapper.h index ad50a8b0d..5d2b4f1cb 100644 --- a/legion/include/taco_mapper.h +++ b/legion/include/taco_mapper.h @@ -44,6 +44,11 @@ class TACOMapper : public Legion::Mapping::DefaultMapper { const Legion::Task& task, TaskOptions& output) override; + void map_task_impl(const Legion::Mapping::MapperContext ctx, + const Legion::Task &task, + const MapTaskInput &input, + MapTaskOutput &output); + void map_task(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, const MapTaskInput &input, @@ -221,6 +226,8 @@ class TACOMapper : public Legion::Mapping::DefaultMapper { // distalReductionConstraintCache is similar to distalLayoutConstraintCache for reduction constraints. std::map, Legion::LayoutConstraintID> distalReductionConstraintCache; + std::map, std::vector> reductionInstanceCache; + // TODO (rohany): It may end up being necessary that we need to explicitly map // regions for placement tasks. If so, Manolis says the following approach // is the right thing: diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 79ba77ed9..0d6a56edc 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -61,6 +61,7 @@ void top_level_task(const Task* task, const std::vector& regions if (dump) { runtime->fill_field(ctx, A.vals, A.valsParent, FID_VAL, valType(0)); } if (batched) { computeLegionBatched(ctx, runtime, &A, &B, &C, &packBatched, pieces); + 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); } diff --git a/legion/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index cd2796d64..0e9ca3ac4 100644 --- a/legion/spmm/taco-generated.cpp +++ b/legion/spmm/taco-generated.cpp @@ -8,6 +8,7 @@ using namespace Legion; typedef FieldAccessor> AccessorROdouble1; typedef FieldAccessor> AccessorROdouble2; typedef FieldAccessor> AccessorRWdouble2; +typedef ReductionAccessor,true,2,coord_t,Realm::AffineAccessor> AccessorReducedouble2; typedef FieldAccessor> AccessorROint32_t1; typedef FieldAccessor,1,coord_t,Realm::AffineAccessor,1,coord_t>> AccessorRORect_1_1; @@ -19,8 +20,14 @@ struct task_1Args { }; struct task_2Args { + int64_t B2Size; + int32_t gx; + int64_t pointID1; +}; + +struct task_3Args { int64_t A2_dimension; - int64_t B1_dimension; + int64_t B2Size; int64_t C2_dimension; int32_t gx; int64_t jo; @@ -71,9 +78,9 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg } BColoring[(*itr)] = BRect; } - auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_DISJOINT_COMPLETE_KIND); + auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_COMPUTE_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_DISJOINT_COMPLETE_KIND); + auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_COMPUTE_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)); @@ -177,16 +184,73 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* } +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 fposo = task->index_point[0]; + task_2Args* args = (task_2Args*)(task->args); + int64_t B2Size = args->B2Size; + int32_t gx = args->gx; + int64_t pointID1 = args->pointID1; + + 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 ; + + DomainT<1> B2PosDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_pos)); + DomainT<1> B2CrdDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)); + int64_t pB2_begin = B2PosDomain.bounds.lo; + int64_t pB2_end = B2PosDomain.bounds.hi; + int64_t pointID2 = pointID1 * gx + fposo; + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx); + int64_t i_pos = taco_binarySearchBefore(B2_pos_accessor, pB2_begin, pB2_end, fposB); + int64_t i = i_pos; + for (int64_t fposi = 0; fposi < ((B2Size + (gx - 1)) / gx); fposi++) { + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx) + fposi; + if (fposB >= (fposo + 1) * ((B2Size + (gx - 1)) / gx)) + continue; + + if (fposB >= B2Size) + continue; + + int64_t f = B2_crd_accessor[fposB]; + int64_t k = f; + while (!(B2_pos_accessor[i_pos].contains(fposB))) { + i_pos = i_pos + 1; + i = i_pos; + } + int64_t iA = i; + int64_t kC = k; + } +} + partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; + auto A_vals_parent = A->valsParent; 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]; + auto B2_crd_parent = B->indicesParents[1][1]; RegionWrapper B_vals = B->vals; + auto B_vals_parent = B->valsParent; IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; int C2_dimension = C->dims[1]; + RegionWrapper C_vals = C->vals; + auto C_vals_parent = C->valsParent; IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; auto computePartitions = partitionPackForcomputeLegionBatched(); @@ -197,61 +261,93 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); - auto ioIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); - DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(ioIndexSpace)); + auto fposoIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(fposoIndexSpace)); + DomainT<1> B2_crd_domain = runtime->get_index_space_domain(ctx, B2_crd.get_index_space()); + DomainPointColoring B2_crd_coloring = DomainPointColoring(); 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<1> itr = PointInDomainIterator<1>(domain); itr.valid(); itr++) { - int64_t io = (*itr)[0]; - Point<2> AStart = Point<2>((io * ((B1_dimension + (gx - 1)) / gx)), (jo * 4)); - 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 * 4 + 3), ADomain.hi()[1])); + int64_t fposo = (*itr)[0]; + Point<1> B2CrdStart = Point<1>((fposo * ((B2Size + (gx - 1)) / gx))); + Point<1> B2CrdEnd = Point<1>(TACO_MIN((fposo * ((B2Size + (gx - 1)) / gx) + ((B2Size + (gx - 1)) / gx - 1)), B2_crd_domain.bounds.hi[0])); + Rect<1> B2CrdRect = Rect<1>(B2CrdStart, B2CrdEnd); + if (!B2_crd_domain.contains(B2CrdRect.lo) || !B2_crd_domain.contains(B2CrdRect.hi)) { + B2CrdRect = B2CrdRect.make_empty(); + } + B2_crd_coloring[(*itr)] = B2CrdRect; + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 4)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension), ADomain.hi()[0]), TACO_MIN((jo * 4 + 3), 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(); + Point<2> CStart = Point<2>(0, (jo * 4)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1), CDomain.hi()[0]), TACO_MIN((jo * 4 + 3), CDomain.hi()[1])); + Rect<2> CRect = Rect<2>(CStart, CEnd); + if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { + CRect = CRect.make_empty(); } - BColoring[(*itr)] = BRect; + 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), pointID1); - auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_DISJOINT_COMPLETE_KIND); - LogicalPartition posPartB2 = copyPartition(ctx, runtime, B_dense_run_0_Partition, B2_pos, pointID1); - LogicalPartition crdPartB2 = runtime->get_logical_partition(ctx, B2_crd, RectCompressedPosPartitionDownwards::apply( + IndexPartition B2_crd_index_part = runtime->create_index_partition(ctx, B2_crd.get_index_space(), domain, B2_crd_coloring, LEGION_COMPUTE_KIND, pointID1); + LogicalPartition B2_crd_part = runtime->get_logical_partition(ctx, B2_crd, B2_crd_index_part); + IndexPartition posSparsePartB2 = runtime->create_partition_by_preimage_range( ctx, - runtime, - B2_crd.get_index_space(), - posPartB2, + B2_crd_index_part, + B2_pos, B2_pos_parent, FID_RECT_1, - pointID1 - )); - auto B_vals_partition = copyPartition(ctx, runtime, crdPartB2, get_logical_region(B_vals), pointID1); + runtime->get_index_partition_color_space_name(ctx, B2_crd_index_part), + LEGION_ALIASED_INCOMPLETE_KIND + ); + IndexPartition posIndexPartB2 = densifyPartition(ctx, runtime, get_index_space(B2_pos), posSparsePartB2, pointID1); + LogicalPartition posPartB2 = runtime->get_logical_partition(ctx, B2_pos, posIndexPartB2); + LogicalPartition BValsLogicalPart = copyPartition(ctx, runtime, B2_crd_part, B_vals, pointID1); + IndexPartition BDenseRun0Partition = copyPartition(ctx, runtime, posPartB2, B_dense_run_0); + IndexPartition ADenseRun0Partition = AffineProjection(0).addOverrides(1).apply(ctx, runtime, BDenseRun0Partition, A_dense_run_0, AColoring); + auto A_vals_partition = copyPartition(ctx, runtime, ADenseRun0Partition, get_logical_region(A_vals), pointID1); + IndexPartition CDenseRun0Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND); + auto C_vals_partition = copyPartition(ctx, runtime, CDenseRun0Partition, get_logical_region(C_vals), pointID1); + task_2Args taskArgsRaw2; + taskArgsRaw2.B2Size = B2Size; + taskArgsRaw2.gx = gx; + taskArgsRaw2.pointID1 = pointID1; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); + IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(A_vals_partition, 0, READ_ONLY, EXCLUSIVE, A_vals_parent, Mapping::DefaultMapper::VIRTUAL_MAP).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(posPartB2, 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1)); + launcher.add_region_requirement(RegionRequirement(B2_crd_part, 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD)); + launcher.add_region_requirement(RegionRequirement(BValsLogicalPart, 0, READ_ONLY, EXCLUSIVE, B_vals_parent, Mapping::DefaultMapper::VIRTUAL_MAP).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(C_vals_partition, 0, READ_ONLY, EXCLUSIVE, C_vals_parent, Mapping::DefaultMapper::VIRTUAL_MAP).add_field(FID_VAL)); + launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; + runtime->execute_index_space(ctx, launcher); + + 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.APartition.denseLevelRunPartitions[0] = ADenseRun0Partition; 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.BPartition.indicesPartitions[1].push_back(B2_crd_part); + computePartitions.BPartition.valsPartition = BValsLogicalPart; + computePartitions.BPartition.denseLevelRunPartitions[0] = BDenseRun0Partition; + computePartitions.CPartition.indicesPartitions = std::vector>(2); + computePartitions.CPartition.denseLevelRunPartitions = std::vector(2); + computePartitions.CPartition.valsPartition = C_vals_partition; + computePartitions.CPartition.denseLevelRunPartitions[0] = CDenseRun0Partition; } return computePartitions; } -void task_2(const Task* task, const std::vector& regions, Context ctx, Runtime* runtime) { +void task_3(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]; @@ -263,10 +359,10 @@ void task_2(const Task* task, const std::vector& regions, Contex PhysicalRegion C_vals = regions[4]; LogicalRegion C_vals_parent = regions[4].get_logical_region(); - int64_t io = task->index_point[0]; - task_2Args* args = (task_2Args*)(task->args); + int64_t fposo = task->index_point[0]; + task_3Args* args = (task_3Args*)(task->args); int64_t A2_dimension = args->A2_dimension; - int64_t B1_dimension = args->B1_dimension; + int64_t B2Size = args->B2Size; int64_t C2_dimension = args->C2_dimension; int32_t gx = args->gx; int64_t jo = args->jo; @@ -274,36 +370,46 @@ void task_2(const Task* task, const std::vector& regions, Contex 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 A_vals_red_accessor = createAccessor(A_vals, FID_VAL, LEGION_REDOP_SUM_FLOAT64); auto B2_pos_accessor = createAccessor(B2_pos, FID_RECT_1); auto B2_crd_accessor = createAccessor(B2_crd, FID_COORD); - int64_t pointID2 = pointID1 * gx + io; - #pragma omp parallel for schedule(dynamic, 128) - 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) + if (runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).empty()) + return ; + + DomainT<1> B2PosDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_pos)); + DomainT<1> B2CrdDomain = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)); + int64_t pB2_begin = B2PosDomain.bounds.lo; + int64_t pB2_end = B2PosDomain.bounds.hi; + int64_t pointID2 = pointID1 * gx + fposo; + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx); + int64_t i_pos = taco_binarySearchBefore(B2_pos_accessor, pB2_begin, pB2_end, fposB); + int64_t i = i_pos; + for (int64_t fposi = 0; fposi < ((B2Size + (gx - 1)) / gx); fposi++) { + int64_t fposB = fposo * ((B2Size + (gx - 1)) / gx) + fposi; + if (fposB >= (fposo + 1) * ((B2Size + (gx - 1)) / gx)) continue; - if (i >= (io + 1) * ((B1_dimension + (gx - 1)) / gx)) + if (fposB >= B2Size) continue; - int64_t pointID3 = pointID2 * ((B1_dimension + (gx - 1)) / gx) + ii; + int64_t f = B2_crd_accessor[fposB]; + int64_t k = f; + while (!(B2_pos_accessor[i_pos].contains(fposB))) { + i_pos = i_pos + 1; + i = i_pos; + } int64_t iA = i; - int64_t iB = i; - 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; - for (int64_t ji = 0; ji < 4; ji++) { - int64_t j = jo * 4 + ji; - if (j >= C2_dimension) - continue; - - int64_t pointID4 = pointID3 * 4 + ji; - int64_t jA = iA * A2_dimension + j; - 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)]; - } + int64_t kC = k; + for (int64_t ji = 0; ji < 4; ji++) { + int64_t j = jo * 4 + ji; + if (j >= C2_dimension) + continue; + + int64_t pointID3 = pointID2 * 4 + ji; + int64_t jA = iA * A2_dimension + j; + int64_t jC = kC * C2_dimension + j; + A_vals_red_accessor[Point<2>(i, j)] <<= B_vals_ro_accessor[Point<1>(fposB)] * C_vals_ro_accessor[Point<2>(k, j)]; } } } @@ -312,7 +418,6 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int A2_dimension = A->dims[1]; RegionWrapper A_vals = A->vals; auto A_vals_parent = A->valsParent; - 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]; @@ -329,22 +434,22 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); - auto ioIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); - DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(ioIndexSpace)); - task_2Args taskArgsRaw2; - taskArgsRaw2.A2_dimension = A2_dimension; - taskArgsRaw2.B1_dimension = B1_dimension; - taskArgsRaw2.C2_dimension = C2_dimension; - taskArgsRaw2.gx = gx; - taskArgsRaw2.jo = jo; - taskArgsRaw2.pointID1 = pointID1; - TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); - IndexLauncher launcher = IndexLauncher(taskID(2), domain, taskArgs, ArgumentMap()); - launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(A_vals), pointID1), 0, READ_WRITE, EXCLUSIVE, A_vals_parent).add_field(FID_VAL)); + auto fposoIndexSpace = runtime->create_index_space(ctx, Rect<1>(lowerBound, upperBound)); + DomainT<1> domain = runtime->get_index_space_domain(ctx, IndexSpaceT<1>(fposoIndexSpace)); + task_3Args taskArgsRaw3; + taskArgsRaw3.A2_dimension = A2_dimension; + taskArgsRaw3.B2Size = B2Size; + taskArgsRaw3.C2_dimension = C2_dimension; + taskArgsRaw3.gx = gx; + taskArgsRaw3.jo = jo; + taskArgsRaw3.pointID1 = pointID1; + TaskArgument taskArgs = TaskArgument(&taskArgsRaw3, sizeof(task_3Args)); + IndexLauncher launcher = IndexLauncher(taskID(3), domain, taskArgs, ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(A_vals), pointID1), 0, LEGION_REDOP_SUM_FLOAT64, LEGION_SIMULTANEOUS, A_vals_parent).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_pos), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_pos_parent)).add_field(FID_RECT_1)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_crd), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 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(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; runtime->execute_index_space(ctx, launcher); @@ -363,4 +468,10 @@ void registerTacoTasks() { registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_2"); } + { + TaskVariantRegistrar registrar(taskID(3), "task_3"); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); + registrar.set_leaf(); + Runtime::preregister_task_variant(registrar, "task_3"); + } } diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index f5435613f..6b8313dd0 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -247,6 +247,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co 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 CColoring = DomainPointColoring(); for (PointInDomainIterator<1> itr = PointInDomainIterator<1>(domain); itr.valid(); itr++) { int64_t fposo = (*itr)[0]; @@ -257,6 +258,13 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 4)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 4 + 3),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<2> CStart = Point<2>(0, (jo * 4)); Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 4 + 3),CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); @@ -287,7 +295,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co LogicalPartition posPartB2 = runtime->get_logical_partition(ctx, B2_pos, posIndexPartB2); LogicalPartition BValsLogicalPart = copyPartition(ctx, runtime, B2_crd_part, B_vals, pointID1); IndexPartition BDenseRun0Partition = copyPartition(ctx, runtime, posPartB2, B_dense_run_0); - IndexPartition ADenseRun0Partition = AffineProjection(0).apply(ctx, runtime, BDenseRun0Partition, A_dense_run_0); + IndexPartition ADenseRun0Partition = AffineProjection(0).addOverrides(1).apply(ctx, runtime, BDenseRun0Partition, A_dense_run_0, AColoring); auto A_vals_partition = copyPartition(ctx, runtime, ADenseRun0Partition, get_logical_region(A_vals), pointID1); IndexPartition CDenseRun0Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND); auto C_vals_partition = copyPartition(ctx, runtime, CDenseRun0Partition, get_logical_region(C_vals), pointID1); diff --git a/legion/spmm/taco-generated.h b/legion/spmm/taco-generated.h index ad9a79f7b..ad2fd96f6 100644 --- a/legion/spmm/taco-generated.h +++ b/legion/spmm/taco-generated.h @@ -11,6 +11,7 @@ struct partitionPackForcomputeLegion { struct partitionPackForcomputeLegionBatched { LegionTensorPartition APartition; LegionTensorPartition BPartition; + LegionTensorPartition CPartition; }; @@ -19,6 +20,7 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegion* partitionPack, int32_t gx); + partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); diff --git a/legion/src/taco_legion_header.cpp b/legion/src/taco_legion_header.cpp index 434fc8568..48ae0b252 100644 --- a/legion/src/taco_legion_header.cpp +++ b/legion/src/taco_legion_header.cpp @@ -293,7 +293,7 @@ int AffineProjection::operator[] (size_t i) const { } Legion::IndexPartition AffineProjection::apply(Legion::Context ctx, Runtime *runtime, Legion::IndexPartition part, - Legion::IndexSpace ispace, Color color) { + Legion::IndexSpace ispace, DomainPointColoring coloringOverride, Color color) { DomainPointColoring col; auto colorSpace = runtime->get_index_partition_color_space(ctx, part); auto colorSpaceName = runtime->get_index_partition_color_space_name(ctx, part); @@ -303,7 +303,19 @@ Legion::IndexPartition AffineProjection::apply(Legion::Context ctx, Runtime *run auto subspace = runtime->get_index_subspace(ctx, part, Color(*itr)); auto subspaceDom = runtime->get_index_space_domain(ctx, subspace); taco_iassert(subspaceDom.dense()); - auto projected = Domain(this->apply(subspaceDom.lo(), outputDomain.lo()), this->apply(subspaceDom.hi(), outputDomain.hi())); + auto lo = this->apply(subspaceDom.lo(), outputDomain.lo()); + auto hi = this->apply(subspaceDom.hi(), outputDomain.hi()); + + // TODO (rohany): Comment this overriding business. + auto it = coloringOverride.find(*itr); + if (it != coloringOverride.end()) { + for (auto idx : this->overrides) { + lo[idx] = it->second.lo()[idx]; + hi[idx] = it->second.hi()[idx]; + } + } + + auto projected = Domain(lo, hi); col[*itr] = projected; } return runtime->create_partition_by_domain(ctx, ispace, col, colorSpaceName, true /* perform_intersections */, LEGION_COMPUTE_KIND, color); diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index e157646a2..431be2e17 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -15,16 +15,19 @@ Realm::Logger logTacoMapper("tacoMapper"); void register_taco_mapper(Machine machine, Runtime *runtime, const std::set &local_procs) { // If we're supposed to backpressure task executions, then we need to only // have a single mapper per node. Otherwise, we can use a mapper per processor. - bool backpressure = false; + bool backpressure = false, oneMapperPerNode = false; auto args = Legion::Runtime::get_input_args(); for (int i = 1; i < args.argc; i++) { if (strcmp(args.argv[i], "-tm:enable_backpressure") == 0) { backpressure = true; break; } + if (strcmp(args.argv[i], "-tm:one_mapper_per_node") == 0) { + oneMapperPerNode = true; + } } - if (backpressure) { + if (backpressure || oneMapperPerNode) { auto proc = *local_procs.begin(); #ifdef TACO_USE_LOGGING_MAPPER runtime->replace_default_mapper(new Mapping::LoggingWrapper(new TACOMapper(runtime->get_mapper_runtime(), machine, proc, TACOMapperName)), Processor::NO_PROC); @@ -141,11 +144,268 @@ void TACOMapper::select_task_options(const Legion::Mapping::MapperContext ctx, } } +void TACOMapper::map_task_impl(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, + const MapTaskInput &input, MapTaskOutput &output) { + Processor::Kind target_kind = task.target_proc.kind(); + // Get the variant that we are going to use to map this task + VariantInfo chosen = default_find_preferred_variant(task, ctx, + true/*needs tight bound*/, true/*cache*/, target_kind); + output.chosen_variant = chosen.variant; + output.task_priority = default_policy_select_task_priority(ctx, task); + output.postmap_task = false; + // Figure out our target processors + default_policy_select_target_processors(ctx, task, output.target_procs); + Processor target_proc = output.target_procs[0]; + // See if we have an inner variant, if we do virtually map all the regions + // We don't even both caching these since they are so simple + if (chosen.is_inner) { + // Check to see if we have any relaxed coherence modes in which + // case we can no longer do virtual mappings so we'll fall through + bool has_relaxed_coherence = false; + for (unsigned idx = 0; idx < task.regions.size(); idx++) { + if (task.regions[idx].prop != LEGION_EXCLUSIVE) { + has_relaxed_coherence = true; + break; + } + } + if (!has_relaxed_coherence) { + std::vector reduction_indexes; + for (unsigned idx = 0; idx < task.regions.size(); idx++) { + // As long as this isn't a reduction-only region requirement + // we will do a virtual mapping, for reduction-only instances + // we will actually make a physical instance because the runtime + // doesn't allow virtual mappings for reduction-only privileges + if (task.regions[idx].privilege == LEGION_REDUCE) + reduction_indexes.push_back(idx); + else + output.chosen_instances[idx].push_back( + PhysicalInstance::get_virtual_instance()); + } + if (!reduction_indexes.empty()) { + const TaskLayoutConstraintSet &layout_constraints = + runtime->find_task_layout_constraints(ctx, + task.task_id, output.chosen_variant); + for (std::vector::const_iterator it = + reduction_indexes.begin(); it != + reduction_indexes.end(); it++) { + MemoryConstraint mem_constraint = + find_memory_constraint(ctx, task, output.chosen_variant, *it); + Memory target_memory = default_policy_select_target_memory(ctx, + target_proc, + task.regions[*it], + mem_constraint); + std::set copy = task.regions[*it].privilege_fields; + size_t footprint; + if (!default_create_custom_instances(ctx, target_proc, + target_memory, task.regions[*it], *it, copy, + layout_constraints, false/*needs constraint check*/, + output.chosen_instances[*it], &footprint)) { + default_report_failed_instance_creation(task, *it, + target_proc, target_memory, footprint); + } + } + } + return; + } + } + // Should we cache this task? + CachedMappingPolicy cache_policy = + default_policy_select_task_cache_policy(ctx, task); + + // First, let's see if we've cached a result of this task mapping + const unsigned long long task_hash = compute_task_hash(task); + std::pair cache_key(task.task_id, target_proc); + std::map, + std::list >::const_iterator + finder = cached_task_mappings.find(cache_key); + // This flag says whether we need to recheck the field constraints, + // possibly because a new field was allocated in a region, so our old + // cached physical instance(s) is(are) no longer valid + bool needs_field_constraint_check = false; + if (cache_policy == DEFAULT_CACHE_POLICY_ENABLE && finder != cached_task_mappings.end()) { + bool found = false; + // Iterate through and see if we can find one with our variant and hash + for (std::list::const_iterator it = + finder->second.begin(); it != finder->second.end(); it++) { + if ((it->variant == output.chosen_variant) && + (it->task_hash == task_hash)) { + // Have to copy it before we do the external call which + // might invalidate our iterator + output.chosen_instances = it->mapping; + output.output_targets = it->output_targets; + output.output_constraints = it->output_constraints; + found = true; + break; + } + } + if (found) { + // See if we can acquire these instances still + if (runtime->acquire_and_filter_instances(ctx, + output.chosen_instances)) + return; + // We need to check the constraints here because we had a + // prior mapping and it failed, which may be the result + // of a change in the allocated fields of a field space + needs_field_constraint_check = true; + // If some of them were deleted, go back and remove this entry + // Have to renew our iterators since they might have been + // invalidated during the 'acquire_and_filter_instances' call + default_remove_cached_task(ctx, output.chosen_variant, + task_hash, cache_key, output.chosen_instances); + } + } + // We didn't find a cached version of the mapping so we need to + // do a full mapping, we already know what variant we want to use + // so let's use one of the acceleration functions to figure out + // which instances still need to be mapped. + std::vector > missing_fields(task.regions.size()); + runtime->filter_instances(ctx, task, output.chosen_variant, + output.chosen_instances, missing_fields); + // Track which regions have already been mapped + std::vector done_regions(task.regions.size(), false); + if (!input.premapped_regions.empty()) + for (std::vector::const_iterator it = + input.premapped_regions.begin(); it != + input.premapped_regions.end(); it++) + done_regions[*it] = true; + const TaskLayoutConstraintSet &layout_constraints = + runtime->find_task_layout_constraints(ctx, + task.task_id, output.chosen_variant); + // Now we need to go through and make instances for any of our + // regions which do not have space for certain fields + for (unsigned idx = 0; idx < task.regions.size(); idx++) { + if (done_regions[idx]) + continue; + // Skip any empty regions + if ((task.regions[idx].privilege == LEGION_NO_ACCESS) || + (task.regions[idx].privilege_fields.empty()) || + missing_fields[idx].empty()) + continue; + // See if this is a reduction + MemoryConstraint mem_constraint = + find_memory_constraint(ctx, task, output.chosen_variant, idx); + Memory target_memory = default_policy_select_target_memory(ctx, + target_proc, + task.regions[idx], + mem_constraint); + if (task.regions[idx].privilege == LEGION_REDUCE) { + // std::cout << "Considering region requirement: " << Utilities::to_string(runtime, ctx, task.regions[idx], idx) << std::endl; + // auto key = std::make_tuple(task.regions[idx].region, size_t(idx), target_memory); + // auto key = std::make_tuple(task.index_point[0], size_t(idx), target_memory); + auto key = std::make_tuple(task_hash, size_t(idx), target_memory); + // std::cout << "KEY: " << std::get<0>(key) << " " << std::get<1>(key) << " " << std::get<2>(key) << std::endl; + auto it = this->reductionInstanceCache.find(key); + if (it != this->reductionInstanceCache.end()) { + std::vector instances = it->second; + // std::cout << "Potential reduction instances: "; + // for (auto inst : instances) { + // std::cout << Utilities::to_string(runtime, ctx, inst) << " "; + // } + // std::cout << std::endl; + if (instances.size() > 0 && this->runtime->acquire_and_filter_instances(ctx, instances)) { + output.chosen_instances[idx] = instances; + this->reductionInstanceCache[key] = instances; + // std::cout << "Reusing instance for regionReq: "; + // for (auto inst : instances) { + // std::cout << Utilities::to_string(runtime, ctx, inst) << " "; + // } + // std::cout << std::endl; + continue; + } else { + // std::cout << "Clearing reduction instance cache?" << std::endl; + this->reductionInstanceCache[key].clear(); + } + } + } + // if (task.regions[idx].privilege == LEGION_REDUCE) { + // size_t footprint; + // if (!default_create_custom_instances(ctx, target_proc, + // target_memory, task.regions[idx], idx, missing_fields[idx], + // layout_constraints, needs_field_constraint_check, + // output.chosen_instances[idx], &footprint)) { + // default_report_failed_instance_creation(task, idx, + // target_proc, target_memory, footprint); + // } + // continue; + // } + // Did the application request a virtual mapping for this requirement? + if ((task.regions[idx].tag & DefaultMapper::VIRTUAL_MAP) != 0 && !(task.regions[idx].privilege == LEGION_REDUCE)) { + PhysicalInstance virt_inst = PhysicalInstance::get_virtual_instance(); + output.chosen_instances[idx].push_back(virt_inst); + continue; + } + // Check to see if any of the valid instances satisfy this requirement + { + std::vector valid_instances; + + for (std::vector::const_iterator + it = input.valid_instances[idx].begin(), + ie = input.valid_instances[idx].end(); it != ie; ++it) { + if (it->get_location() == target_memory) + valid_instances.push_back(*it); + } + std::set valid_missing_fields; + runtime->filter_instances(ctx, task, idx, output.chosen_variant, + valid_instances, valid_missing_fields); + +#ifndef NDEBUG + bool check = +#endif + runtime->acquire_and_filter_instances(ctx, valid_instances); + assert(check); + + output.chosen_instances[idx] = valid_instances; + missing_fields[idx] = valid_missing_fields; + + if (missing_fields[idx].empty()) { + continue; + } + } + // Otherwise make normal instances for the given region + size_t footprint; + if (!default_create_custom_instances(ctx, target_proc, + target_memory, task.regions[idx], idx, missing_fields[idx], + layout_constraints, needs_field_constraint_check, + output.chosen_instances[idx], &footprint)) { + default_report_failed_instance_creation(task, idx, + target_proc, target_memory, footprint); + } + if (task.regions[idx].privilege == LEGION_REDUCE) { + // auto key = std::make_tuple(task.regions[idx].region, size_t(idx), target_memory); + // auto key = std::make_tuple(task.index_point[0], size_t(idx), target_memory); + auto key = std::make_tuple(task_hash, size_t(idx), target_memory); + // auto key = std::make_tuple(task_hash, size_t(idx), target_memory); + this->reductionInstanceCache[key] = output.chosen_instances[idx]; + } + } + + // Finally we set a target memory for output instances + Memory target_memory = + default_policy_select_output_target(ctx, task.target_proc); + for (unsigned i = 0; i < task.output_regions.size(); ++i) { + output.output_targets[i] = target_memory; + default_policy_select_output_constraints( + task, output.output_constraints[i], task.output_regions[i]); + } + + if (cache_policy == DEFAULT_CACHE_POLICY_ENABLE && false) { + // Now that we are done, let's cache the result so we can use it later + std::list &map_list = cached_task_mappings[cache_key]; + map_list.push_back(CachedTaskMapping()); + CachedTaskMapping &cached_result = map_list.back(); + cached_result.task_hash = task_hash; + cached_result.variant = output.chosen_variant; + cached_result.mapping = output.chosen_instances; + cached_result.output_targets = output.output_targets; + cached_result.output_constraints = output.output_constraints; + } +} + void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, const MapTaskInput &input, MapTaskOutput &output) { - DefaultMapper::map_task(ctx, task, input, output); + this->map_task_impl(ctx, task, input, output); // If the tag is marked for untracked valid regions, then mark all of its // read only regions as up for collection. if ((task.tag & UNTRACK_VALID_REGIONS) != 0 && this->untrackValidRegions) { diff --git a/src/lower/lowerer_impl.cpp b/src/lower/lowerer_impl.cpp index b6a54a595..f787b9cb8 100644 --- a/src/lower/lowerer_impl.cpp +++ b/src/lower/lowerer_impl.cpp @@ -2186,6 +2186,7 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, std::map extraColorings; + std::map> partitionedByOuter; std::vector colorings; std::vector initColorings; for (auto &t : forall.getTransfers()) { @@ -2215,19 +2216,21 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, // std::set considering = this->varsInScope[this->curDistVar]; std::set considering = this->definedIndexVars; considering.erase(this->curDistVar); - for (auto ivar : t.getAccess().getIndexVars()) { + for (size_t i = 0; i < t.getAccess().getIndexVars().size(); i++) { + auto ivar = t.getAccess().getIndexVars()[i]; if (this->anyParentInSet(ivar, considering)) { isPartitionedByOuterVars = true; + partitionedByOuter[tv].push_back(i); } } - if (isPartitionedByOuterVars && !affineProjection.defined()) { + if (isPartitionedByOuterVars) { partitioningStmts.push_back(initColorings[idx]); std::set fullyRepl; coloringLoopBody.push_back(ir::Block::make(this->createDomainPointColorings(forall, domainIter, domains, fullyRepl, colorings, tv))); extraColorings[tv] = colorings[idx]; } // TODO (rohany): As described below, we aren't going to be able to compute a tight bound - // on the subspaces here. + // on the subspaces here. I think I know what to do here now? // taco_iassert(!(isPartitionedByOuterVars && affineProjection.defined())); } auto coloringLoop = ir::For::make( @@ -2426,17 +2429,28 @@ Stmt LowererImpl::lowerForallDimension(Forall forall, auto tvExpr = this->tensorVars[tv]; auto tvAccess = this->tensorVarToAccess.at(tv); + auto affineProjection = this->constructAffineProjection(posAccess, tvAccess); + ir::Expr createDenseRunPart; auto firstDenseRun = ir::GetProperty::makeDenseLevelRun(tvExpr, 0); - if (util::contains(extraColorings, tv)) { + if (util::contains(extraColorings, tv) && !affineProjection.defined()) { createDenseRunPart = ir::Call::make("runtime->create_index_partition", {ctx, firstDenseRun, domain, extraColorings[tv], computePart}, Auto); } else { // Attempt to create different projections from the pos tensor // into the target tensor. - auto affineProjection = this->constructAffineProjection(posAccess, tvAccess); auto sparseGatherProjection = this->constructSparseGatherProjection(posAccess, tvAccess, posMode); if (affineProjection.defined()) { - createDenseRunPart = ir::MethodCall::make(affineProjection, "apply", {ctx, runtime, denseRunPartition, firstDenseRun}, false /* deref */, IndexPartition); + std::vector projectionArgs = {ctx, runtime, denseRunPartition, firstDenseRun}; + if (util::contains(partitionedByOuter, tv)) { + auto overrides = partitionedByOuter[tv]; + std::vector overrideExprs; + for (auto it : overrides) { + overrideExprs.push_back(Expr(it)); + } + affineProjection = ir::MethodCall::make(affineProjection, "addOverrides", overrideExprs, false /* deref */, Auto); + projectionArgs.push_back(extraColorings[tv]); + } + createDenseRunPart = ir::MethodCall::make(affineProjection, "apply", projectionArgs, false /* deref */, IndexPartition); } else if (sparseGatherProjection.defined() && !isFusedPos) { // We can't apply a SparseGatherProjection if the position space loop is fused. // TODO (rohany): This is still hardcoded for the RectCompressedModeFormat. @@ -6084,7 +6098,7 @@ std::vector LowererImpl::createIndexPartitions( auto part = ir::Var::make(tv.getName() + "_dense_run_0_Partition", Auto); // TODO (rohany): Make this LEGION_COMPUTE_KIND since it is happening // off of the critical path. - auto partKind = disjointPart; + auto partKind = computePart; // Figure out how many axes of the tensor are not being partitioned in order // to figure out how many axes of the tensor are being partitioned. If // the tensor is being partitioned in as many ways as the target loop is diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp index 97d625625..a73a6cba7 100644 --- a/test/tests-distributed.cpp +++ b/test/tests-distributed.cpp @@ -1849,10 +1849,15 @@ TEST(distributed, legionSpMM) { const int BATCH_SIZE = 4; // Schedule to conserve memory (i.e. not replicate C onto all nodes). auto cpuCons = stmt.split(j, jo, ji, BATCH_SIZE) + // TODO (rohany): Let's switch this back to a pos split to check the reduction instance stuff. .reorder({jo, i, k, ji}) - .distribute({i}, {io}, {ii}, Grid(gx)) - .parallelize(ii, ParallelUnit::CPUThread, OutputRaceStrategy::NoRaces) - .communicate({A(i, j), B(i, k), C(k, j)}, io) + .fuse(i, k, f) + .pos(f, fpos, B(i, k)) + .distribute({fpos}, {fposo}, {fposi}, Grid(gx)) + .communicate({A(i, j), B(i, k), C(k, j)}, fposi) + // .distribute({i}, {io}, {ii}, Grid(gx)) + // .parallelize(ii, ParallelUnit::CPUThread, OutputRaceStrategy::NoRaces) + // .communicate({A(i, j), B(i, k), C(k, j)}, io) ; const int BATCHED_NNZ_PER_THREAD = 8; From a2fc54c4a704df6309566ce1294c64affe338c30 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Sat, 12 Feb 2022 00:43:53 -0800 Subject: [PATCH 03/12] fixup --- legion/spmm/taco-generated.cu | 26 +++++++++++++------------- legion/src/taco_mapper.cpp | 2 +- src/codegen/codegen.cpp | 2 +- src/codegen/codegen_legion.cpp | 3 +-- test/tests-distributed.cpp | 2 +- 5 files changed, 17 insertions(+), 18 deletions(-) diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index 6b8313dd0..a72c380a4 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -192,7 +192,7 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* 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]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; auto C_vals_parent = C->valsParent; @@ -221,14 +221,14 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { 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]; + size_t B1_dimension = B->dims[0]; + size_t 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]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; @@ -236,7 +236,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -258,15 +258,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 4)); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 4 + 3),ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 8)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 8 + 7),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<2> CStart = Point<2>(0, (jo * 4)); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 4 + 3),CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, (jo * 8)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 8 + 7),CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { CRect = CRect.make_empty(); @@ -355,8 +355,8 @@ void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_b i_pos = i_pos + 1; i = i_pos; } - for (int64_t ji = 0; ji < 4; ji++) { - int64_t j = jo * 4 + ji; + for (int64_t ji = 0; ji < 8; ji++) { + int64_t j = jo * 8 + ji; if (j >= C2_dimension) break; @@ -423,13 +423,13 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT auto B2_crd_parent = B->indicesParents[1][1]; RegionWrapper B_vals = B->vals; auto B_vals_parent = B->valsParent; - int C2_dimension = C->dims[1]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; auto C_vals_parent = C->valsParent; int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index 431be2e17..19fb0d397 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -411,7 +411,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); } } diff --git a/src/codegen/codegen.cpp b/src/codegen/codegen.cpp index b89b68c26..b65256f3a 100644 --- a/src/codegen/codegen.cpp +++ b/src/codegen/codegen.cpp @@ -236,7 +236,7 @@ string CodeGen::printTensorProperty(string varname, const GetProperty* op, bool // for a Fixed level, ptr is an int // all others are int* if (op->property == TensorProperty::Dimension) { - tp = "int" + star; + tp = "size_t" + star; ret << tp << " " << varname; } else { taco_iassert(op->property == TensorProperty::Indices); diff --git a/src/codegen/codegen_legion.cpp b/src/codegen/codegen_legion.cpp index aded6a0b7..06e246e14 100644 --- a/src/codegen/codegen_legion.cpp +++ b/src/codegen/codegen_legion.cpp @@ -16,8 +16,7 @@ std::string CodegenLegion::unpackTensorProperty(std::string varname, const GetPr // tp = "int"; // ret << tp << " " << varname << " = runtime->get_index_space_domain(get_index_space(" << tensor->name << // ")).hi()[" << op->mode << "] + 1;\n"; - // TODO (rohany): Does this need to be an int64? - ret << "int " << varname << " = " << tensor->name << "->dims[" << op->mode << "];\n"; + ret << "size_t " << varname << " = " << tensor->name << "->dims[" << op->mode << "];\n"; } else if (op->property == TensorProperty::IndexSpace) { tp = "auto"; ret << tp << " " << varname << " = get_index_space(" << tensor->name << ");\n"; diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp index a73a6cba7..4ae97dd6e 100644 --- a/test/tests-distributed.cpp +++ b/test/tests-distributed.cpp @@ -1846,7 +1846,7 @@ TEST(distributed, legionSpMM) { auto loweredCPU = lowerLegionSeparatePartitionCompute(cpuStmt, "computeLegion", false /* waitOnFutureMap */); auto loweredGPU = lowerLegionSeparatePartitionCompute(gpuStmt, "computeLegion", false /* waitOnFutureMap */); - const int BATCH_SIZE = 4; + const int BATCH_SIZE = 8; // Schedule to conserve memory (i.e. not replicate C onto all nodes). auto cpuCons = stmt.split(j, jo, ji, BATCH_SIZE) // TODO (rohany): Let's switch this back to a pos split to check the reduction instance stuff. From 7b913bf2d6c90e304f473c77d7ec360b520646ba Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Sat, 12 Feb 2022 00:57:49 -0800 Subject: [PATCH 04/12] fixup --- legion/src/taco_mapper.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index 19fb0d397..32ac37df6 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -411,6 +411,11 @@ 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]; + // TODO (rohany): I only set i == 4 here so that the C region in the SpMM computation + // gets collected instead of any of the other regions, which forces the regions to get + // loaded from disks intead of remaining in FB memory. I need to investigate what + // happens if I create long lasting instances with an explicit dummy read and then + // run with untraked all. I don't think that helps here too... if (rg.privilege == READ_ONLY && i == 4) { output.untracked_valid_regions.insert(i); } From 3c4075a714978190322e071af4c77190149e15c2 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Mon, 14 Feb 2022 10:30:08 -0800 Subject: [PATCH 05/12] fixup --- legion/spmm/taco-generated.cpp | 44 +++++++++---------- legion/spmm/taco-generated.cu | 19 +++++---- legion/src/taco_mapper.cpp | 78 ++++++++++++++++++++++++++++++++-- 3 files changed, 107 insertions(+), 34 deletions(-) diff --git a/legion/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index 0e9ca3ac4..c0c06cc01 100644 --- a/legion/spmm/taco-generated.cpp +++ b/legion/spmm/taco-generated.cpp @@ -38,13 +38,13 @@ struct task_3Args { partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { RegionWrapper A_vals = A->vals; IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; - int B1_dimension = B->dims[0]; + size_t 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]; RegionWrapper B_vals = B->vals; IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; - int C2_dimension = C->dims[1]; + size_t C2_dimension = C->dims[1]; IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; auto computePartitions = partitionPackForcomputeLegion(); @@ -150,14 +150,14 @@ void task_1(const Task* task, const std::vector& regions, Contex } void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegion* partitionPack, int32_t gx) { - int A2_dimension = A->dims[1]; + size_t A2_dimension = A->dims[1]; auto A_vals_parent = A->valsParent; - int B1_dimension = B->dims[0]; + size_t 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]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; auto C_vals_parent = C->valsParent; @@ -239,8 +239,8 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co RegionWrapper A_vals = A->vals; auto A_vals_parent = A->valsParent; IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; - int B1_dimension = B->dims[0]; - int B2_dimension = B->dims[1]; + size_t B1_dimension = B->dims[0]; + size_t 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]; @@ -248,7 +248,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co RegionWrapper B_vals = B->vals; auto B_vals_parent = B->valsParent; IndexSpace B_dense_run_0 = B->denseLevelRuns[0]; - int C2_dimension = C->dims[1]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; auto C_vals_parent = C->valsParent; IndexSpace C_dense_run_0 = C->denseLevelRuns[0]; @@ -257,7 +257,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + for (int64_t jo = 0; jo < C2_dimension; jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -279,15 +279,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 4)); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension), ADomain.hi()[0]), TACO_MIN((jo * 4 + 3), ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), jo); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension), ADomain.hi()[0]), TACO_MIN(jo, 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<2> CStart = Point<2>(0, (jo * 4)); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1), CDomain.hi()[0]), TACO_MIN((jo * 4 + 3), CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, jo); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1), CDomain.hi()[0]), TACO_MIN(jo, CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { CRect = CRect.make_empty(); @@ -401,12 +401,12 @@ void task_3(const Task* task, const std::vector& regions, Contex } int64_t iA = i; int64_t kC = k; - for (int64_t ji = 0; ji < 4; ji++) { - int64_t j = jo * 4 + ji; + for (int64_t ji = 0; ji < 1; ji++) { + int64_t j = jo + ji; if (j >= C2_dimension) continue; - int64_t pointID3 = pointID2 * 4 + ji; + int64_t pointID3 = pointID2 * 1 + ji; int64_t jA = iA * A2_dimension + j; int64_t jC = kC * C2_dimension + j; A_vals_red_accessor[Point<2>(i, j)] <<= B_vals_ro_accessor[Point<1>(fposB)] * C_vals_ro_accessor[Point<2>(k, j)]; @@ -415,7 +415,7 @@ void task_3(const Task* task, const std::vector& regions, Contex } void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx) { - int A2_dimension = A->dims[1]; + size_t A2_dimension = A->dims[1]; RegionWrapper A_vals = A->vals; auto A_vals_parent = A->valsParent; RegionWrapper B2_pos = B->indices[1][0]; @@ -424,13 +424,13 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT auto B2_crd_parent = B->indicesParents[1][1]; RegionWrapper B_vals = B->vals; auto B_vals_parent = B->valsParent; - int C2_dimension = C->dims[1]; + size_t C2_dimension = C->dims[1]; RegionWrapper C_vals = C->vals; auto C_vals_parent = C->valsParent; int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 3) / 4); jo++) { + for (int64_t jo = 0; jo < C2_dimension; jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -458,19 +458,19 @@ void computeLegionBatched(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"); } diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index a72c380a4..b5cd4f5c4 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -236,7 +236,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { + for (int64_t jo = 0; jo < C2_dimension; jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -258,15 +258,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 8)); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 8 + 7),ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), jo); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN(jo,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<2> CStart = Point<2>(0, (jo * 8)); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 8 + 7),CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, jo); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN(jo,CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { CRect = CRect.make_empty(); @@ -355,8 +355,8 @@ void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_b i_pos = i_pos + 1; i = i_pos; } - for (int64_t ji = 0; ji < 8; ji++) { - int64_t j = jo * 8 + ji; + for (int64_t ji = 0; ji < 1; ji++) { + int64_t j = jo + ji; if (j >= C2_dimension) break; @@ -429,7 +429,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { + for (int64_t jo = 0; jo < C2_dimension; jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -449,7 +449,8 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; - runtime->execute_index_space(ctx, launcher); + launcher.tag = launcher.tag | TACOMapper::BACKPRESSURE_TASK; + runtime->execute_index_space(ctx, launcher).wait_all_results(); } } diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index 32ac37df6..eb9b88d6e 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -29,6 +29,7 @@ void register_taco_mapper(Machine machine, Runtime *runtime, const std::setreplace_default_mapper(new Mapping::LoggingWrapper(new TACOMapper(runtime->get_mapper_runtime(), machine, proc, TACOMapperName)), Processor::NO_PROC); #else @@ -142,6 +143,11 @@ void TACOMapper::select_task_options(const Legion::Mapping::MapperContext ctx, if (this->multipleShardsPerNode && task.get_depth() == 0) { output.replicate = true; } + + std::cout << "Task {" << task.get_task_name() << ", " << task.get_unique_id() << "} has select task options result: " + << "{ initial_proc: " << output.initial_proc << ", inline_task: " << output.inline_task << ", stealable: " << output.stealable + << ", map_locally: " << output.map_locally << ", valid_instances: " << output.valid_instances << ", memoize: " << output.memoize + << ", replicate: " << output.replicate << ", parent_priority: " << output.parent_priority << "}" << std::endl; } void TACOMapper::map_task_impl(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, @@ -405,6 +411,14 @@ void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, const MapTaskInput &input, MapTaskOutput &output) { + std::cout << "Mapping task: " << task.get_task_name() << " " << task.get_unique_id() << " " << task.is_index_space << " " << task.index_point << std::endl; + // std::stringstream ss; + // ss << "{"; + // for (size_t idx = 0; idx < task.regions.size(); idx++) { + // ss << Utilities::to_string(runtime, ctx, task.regions[idx], idx) << ", "; + // } + // ss << "}"; + // std::cout << "Task has regs: " << ss.str() << std::endl; this->map_task_impl(ctx, task, input, output); // If the tag is marked for untracked valid regions, then mark all of its // read only regions as up for collection. @@ -423,6 +437,7 @@ void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, } // Mark that we want profiling from this task if we're supposed to backpressure it. if ((task.tag & BACKPRESSURE_TASK) != 0 && this->enableBackpressure) { + std::cout << "Requesting profiling info for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; output.task_prof_requests.add_measurement(); } } @@ -862,11 +877,23 @@ 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, const Task& task, const TaskProfilingInfo& input) { + std::cout << "In profiling result for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; // We should only get profiling responses if we've enabled backpressuring. taco_iassert(this->enableBackpressure); // We should only get profiling responses for tasks that are supposed to be backpressured. @@ -878,7 +905,15 @@ void TACOMapper::report_profiling(const MapperContext ctx, delete prof; // 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]; + // auto& inflight = this->backPressureQueue[task.orig_proc]; + auto& inflight = this->backPressureQueue[task.target_proc]; + std::stringstream ss; + ss << "{ "; + for (auto it : this->backPressureQueue[task.target_proc]) { + ss << it.id << ", "; + } + ss << "}"; + std::cout << "For proc: " << task.target_proc << " inflight " << ss.str() << std::endl; MapperEvent event; // Find this task in the queue. for (auto it = inflight.begin(); it != inflight.end(); it++) { @@ -888,6 +923,9 @@ void TACOMapper::report_profiling(const MapperContext ctx, break; } } + if (!event.exists()) { + std::cout << std::string(task.get_task_name()) << " " << task.index_point << " " << task.get_unique_id() << " " << task.orig_proc << " " << task.target_proc << " " << task.current_proc << " " << inflight.size() << std::endl; + } // Assert that we found a valid event. taco_iassert(event.exists()); // Finally, trigger the event for anyone waiting on it. @@ -899,6 +937,13 @@ void TACOMapper::report_profiling(const MapperContext ctx, void TACOMapper::select_tasks_to_map(const MapperContext ctx, const SelectMappingInput& input, SelectMappingOutput& output) { + { + std::stringstream ss; + for (auto it : input.ready_tasks) { + ss << "{ " << it->get_task_name() << " " << it->get_unique_id() << "}, "; + } + std::cout << "Entering call to select_tasks_to_map, have tasks: " << ss.str() << std::endl; + } if (!this->enableBackpressure) { DefaultMapper::select_tasks_to_map(ctx, input, output); } else { @@ -928,11 +973,22 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, { auto task = *it; bool schedule = true; + if (std::string(task->get_task_name()) == "task_2") { + taco_iassert((task->tag & BACKPRESSURE_TASK) != 0); + } if ((task->tag & BACKPRESSURE_TASK) != 0) { // 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. - auto inflight = this->backPressureQueue[task->orig_proc]; + // auto inflight = this->backPressureQueue[task->orig_proc]; + auto inflight = this->backPressureQueue[task->target_proc]; + std::stringstream ss; + ss << "{ "; + for (auto it : this->backPressureQueue[task->target_proc]) { + ss << it.id << ", "; + } + ss << "}"; + std::cout << "Inflight for task: " << task->get_task_name() << " " << task->get_unique_id() << " " << ss.str() << std::endl; if (inflight.size() == this->maxInFlightTasks) { // We've hit the cap, so we can't schedule any more tasks. schedule = false; @@ -947,16 +1003,32 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, } else { // Otherwise, we can schedule the task. Create a new event // and queue it up on the processor. - this->backPressureQueue[task->orig_proc].push_back({ + std::cout << "Adding event to execution queue " << std::string(task->get_task_name()) << " " << task->index_point << " " << task->orig_proc << " " << task->target_proc << " " << task->current_proc << " " << task->is_index_space << " " << task->index_point << std::endl; + this->backPressureQueue[task->target_proc].push_back({ .id = task->get_unique_id(), .event = this->runtime->create_mapper_event(ctx), .schedTime = schedTime, }); + std::stringstream ss; + ss << "{ "; + for (auto it : this->backPressureQueue[task->target_proc]) { + ss << it.id << ", "; + } + ss << "}"; + std::cout << "target proc " << task->target_proc << " -> " << ss.str() << std::endl; } } // Schedule tasks that are valid and have the target depth. if (schedule && (*it)->get_depth() == max_depth) { + std::cout << "Adding task to map task output: " << (*it)->get_task_name() << " " << (*it)->get_unique_id() << std::endl; + std::stringstream ss; + ss << "{"; + for (size_t idx = 0; idx < task->regions.size(); idx++) { + ss << Utilities::to_string(runtime, ctx, task->regions[idx], idx) << ", "; + } + ss << "}"; + std::cout << "Task has regs: " << ss.str() << std::endl; output.map_tasks.insert(*it); count++; } From 9188c6c542c2b6902d35790c1b64884cd7867aaa Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Tue, 15 Feb 2022 13:35:14 -0800 Subject: [PATCH 06/12] fixup --- legion/spmm/taco-generated.cpp | 27 ++++++++++++++------------- legion/spmm/taco-generated.cu | 19 +++++++++---------- legion/src/taco_mapper.cpp | 10 ++++++++-- 3 files changed, 31 insertions(+), 25 deletions(-) diff --git a/legion/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index c0c06cc01..dde4eca29 100644 --- a/legion/spmm/taco-generated.cpp +++ b/legion/spmm/taco-generated.cpp @@ -257,7 +257,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < C2_dimension; jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -279,15 +279,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), jo); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension), ADomain.hi()[0]), TACO_MIN(jo, ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 8)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension), ADomain.hi()[0]), TACO_MIN((jo * 8 + 7), 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<2> CStart = Point<2>(0, jo); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1), CDomain.hi()[0]), TACO_MIN(jo, CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, (jo * 8)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1), CDomain.hi()[0]), TACO_MIN((jo * 8 + 7), CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { CRect = CRect.make_empty(); @@ -325,7 +325,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co launcher.add_region_requirement(RegionRequirement(BValsLogicalPart, 0, READ_ONLY, EXCLUSIVE, B_vals_parent, Mapping::DefaultMapper::VIRTUAL_MAP).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(C_vals_partition, 0, READ_ONLY, EXCLUSIVE, C_vals_parent, Mapping::DefaultMapper::VIRTUAL_MAP).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; - runtime->execute_index_space(ctx, launcher); + // runtime->execute_index_space(ctx, launcher); computePartitions.APartition.indicesPartitions = std::vector>(2); @@ -401,12 +401,12 @@ void task_3(const Task* task, const std::vector& regions, Contex } int64_t iA = i; int64_t kC = k; - for (int64_t ji = 0; ji < 1; ji++) { - int64_t j = jo + ji; + for (int64_t ji = 0; ji < 8; ji++) { + int64_t j = jo * 8 + ji; if (j >= C2_dimension) continue; - int64_t pointID3 = pointID2 * 1 + ji; + int64_t pointID3 = pointID2 * 8 + ji; int64_t jA = iA * A2_dimension + j; int64_t jC = kC * C2_dimension + j; A_vals_red_accessor[Point<2>(i, j)] <<= B_vals_ro_accessor[Point<1>(fposB)] * C_vals_ro_accessor[Point<2>(k, j)]; @@ -430,7 +430,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < C2_dimension; jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -451,6 +451,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; + launcher.tag = launcher.tag | TACOMapper::BACKPRESSURE_TASK; runtime->execute_index_space(ctx, launcher); } @@ -458,19 +459,19 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT 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"); } { TaskVariantRegistrar registrar(taskID(3), "task_3"); - registrar.add_constraint(ProcessorConstraint(Processor::OMP_PROC)); + registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC)); registrar.set_leaf(); Runtime::preregister_task_variant(registrar, "task_3"); } diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index b5cd4f5c4..a72c380a4 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -236,7 +236,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < C2_dimension; jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -258,15 +258,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), jo); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN(jo,ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 8)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 8 + 7),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<2> CStart = Point<2>(0, jo); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN(jo,CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, (jo * 8)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 8 + 7),CDomain.hi()[1])); Rect<2> CRect = Rect<2>(CStart, CEnd); if (!CDomain.contains(CRect.lo) || !CDomain.contains(CRect.hi)) { CRect = CRect.make_empty(); @@ -355,8 +355,8 @@ void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_b i_pos = i_pos + 1; i = i_pos; } - for (int64_t ji = 0; ji < 1; ji++) { - int64_t j = jo + ji; + for (int64_t ji = 0; ji < 8; ji++) { + int64_t j = jo * 8 + ji; if (j >= C2_dimension) break; @@ -429,7 +429,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < C2_dimension; jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -449,8 +449,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; - launcher.tag = launcher.tag | TACOMapper::BACKPRESSURE_TASK; - runtime->execute_index_space(ctx, launcher).wait_all_results(); + runtime->execute_index_space(ctx, launcher); } } diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index eb9b88d6e..625a51e84 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -859,6 +859,12 @@ 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); + std::cout << "Targets: "; + for (auto it : targets) { + std::cout << it << " " << std::endl; + } + std::cout << std::endl; + bool recurse = ((task.tag & BACKPRESSURE_TASK) != 0) && input.domain.get_volume() > 1; switch (input.domain.get_dim()) { #define BLOCK(DIM) \ case DIM: \ @@ -867,7 +873,7 @@ void TACOMapper::slice_task(const Legion::Mapping::MapperContext ctx, Legion::Point num_blocks = \ default_select_num_blocks(targets.size(), point_space.bounds); \ this->default_decompose_points(point_space, targets, \ - num_blocks, false/*recurse*/, \ + num_blocks, recurse/*recurse*/, \ stealing_enabled, output.slices); \ break; \ } @@ -976,7 +982,7 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, if (std::string(task->get_task_name()) == "task_2") { taco_iassert((task->tag & BACKPRESSURE_TASK) != 0); } - if ((task->tag & BACKPRESSURE_TASK) != 0) { + if ((task->tag & BACKPRESSURE_TASK) != 0 && !(task->is_index_space && task->index_domain.get_volume() > 1)) { // 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. From bf34f50f876fddbc5fcd563bc247c26aed084845 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Wed, 16 Feb 2022 10:00:10 -0800 Subject: [PATCH 07/12] fixup --- legion/spmm/main.cpp | 11 ++++ legion/src/taco_mapper.cpp | 122 ++++++++++++++++++------------------- scripts/spbenchmark.py | 2 +- 3 files changed, 73 insertions(+), 62 deletions(-) diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 0d6a56edc..c093b4854 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -6,7 +6,12 @@ #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; @@ -57,6 +62,11 @@ void top_level_task(const Task* task, const std::vector& regions } auto commLPart = runtime->get_logical_partition(ctx, A.vals, commPart); + // Create a long lasting instance for the C region to force communcation costs. + if (batched) { + tacoFill(ctx, runtime, C.vals, valType(1)); + } + auto avgTime = benchmarkAsyncCallWithWarmup(ctx, runtime, warmup, n, [&]() { if (dump) { runtime->fill_field(ctx, A.vals, A.valsParent, FID_VAL, valType(0)); } if (batched) { @@ -97,6 +107,7 @@ int main(int argc, char** argv) { registerHDF5UtilTasks(); registerTacoTasks(); registerTacoRuntimeLibTasks(); + registerTACOFillTasks(); Runtime::add_registration_callback(register_taco_mapper); Runtime::preregister_sharding_functor(TACOShardingFunctorID, new TACOShardingFunctor()); return Runtime::start(argc, argv); diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index 625a51e84..d9150a594 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -29,7 +29,7 @@ void register_taco_mapper(Machine machine, Runtime *runtime, const std::setreplace_default_mapper(new Mapping::LoggingWrapper(new TACOMapper(runtime->get_mapper_runtime(), machine, proc, TACOMapperName)), Processor::NO_PROC); #else @@ -144,10 +144,10 @@ void TACOMapper::select_task_options(const Legion::Mapping::MapperContext ctx, output.replicate = true; } - std::cout << "Task {" << task.get_task_name() << ", " << task.get_unique_id() << "} has select task options result: " - << "{ initial_proc: " << output.initial_proc << ", inline_task: " << output.inline_task << ", stealable: " << output.stealable - << ", map_locally: " << output.map_locally << ", valid_instances: " << output.valid_instances << ", memoize: " << output.memoize - << ", replicate: " << output.replicate << ", parent_priority: " << output.parent_priority << "}" << std::endl; + // std::cout << "Task {" << task.get_task_name() << ", " << task.get_unique_id() << "} has select task options result: " + // << "{ initial_proc: " << output.initial_proc << ", inline_task: " << output.inline_task << ", stealable: " << output.stealable + // << ", map_locally: " << output.map_locally << ", valid_instances: " << output.valid_instances << ", memoize: " << output.memoize + // << ", replicate: " << output.replicate << ", parent_priority: " << output.parent_priority << "}" << std::endl; } void TACOMapper::map_task_impl(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, @@ -411,7 +411,7 @@ void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, const Legion::Task &task, const MapTaskInput &input, MapTaskOutput &output) { - std::cout << "Mapping task: " << task.get_task_name() << " " << task.get_unique_id() << " " << task.is_index_space << " " << task.index_point << std::endl; + // std::cout << "Mapping task: " << task.get_task_name() << " " << task.get_unique_id() << " " << task.is_index_space << " " << task.index_point << std::endl; // std::stringstream ss; // ss << "{"; // for (size_t idx = 0; idx < task.regions.size(); idx++) { @@ -437,7 +437,7 @@ void TACOMapper::map_task(const Legion::Mapping::MapperContext ctx, } // Mark that we want profiling from this task if we're supposed to backpressure it. if ((task.tag & BACKPRESSURE_TASK) != 0 && this->enableBackpressure) { - std::cout << "Requesting profiling info for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; + // std::cout << "Requesting profiling info for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; output.task_prof_requests.add_measurement(); } } @@ -859,11 +859,11 @@ 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); - std::cout << "Targets: "; - for (auto it : targets) { - std::cout << it << " " << std::endl; - } - std::cout << std::endl; + // std::cout << "Targets: "; + // for (auto it : targets) { + // std::cout << it << " " << std::endl; + // } + // std::cout << std::endl; bool recurse = ((task.tag & BACKPRESSURE_TASK) != 0) && input.domain.get_volume() > 1; switch (input.domain.get_dim()) { #define BLOCK(DIM) \ @@ -883,23 +883,23 @@ 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; - } + // 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, const Task& task, const TaskProfilingInfo& input) { - std::cout << "In profiling result for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; + // std::cout << "In profiling result for task: " << task.get_task_name() << " " << task.get_unique_id() << std::endl; // We should only get profiling responses if we've enabled backpressuring. taco_iassert(this->enableBackpressure); // We should only get profiling responses for tasks that are supposed to be backpressured. @@ -913,13 +913,13 @@ void TACOMapper::report_profiling(const MapperContext ctx, // So, we'll use orig_proc to index into the queue. // auto& inflight = this->backPressureQueue[task.orig_proc]; auto& inflight = this->backPressureQueue[task.target_proc]; - std::stringstream ss; - ss << "{ "; - for (auto it : this->backPressureQueue[task.target_proc]) { - ss << it.id << ", "; - } - ss << "}"; - std::cout << "For proc: " << task.target_proc << " inflight " << ss.str() << std::endl; + // std::stringstream ss; + // ss << "{ "; + // for (auto it : this->backPressureQueue[task.target_proc]) { + // ss << it.id << ", "; + // } + // ss << "}"; + // std::cout << "For proc: " << task.target_proc << " inflight " << ss.str() << std::endl; MapperEvent event; // Find this task in the queue. for (auto it = inflight.begin(); it != inflight.end(); it++) { @@ -929,9 +929,9 @@ void TACOMapper::report_profiling(const MapperContext ctx, break; } } - if (!event.exists()) { - std::cout << std::string(task.get_task_name()) << " " << task.index_point << " " << task.get_unique_id() << " " << task.orig_proc << " " << task.target_proc << " " << task.current_proc << " " << inflight.size() << std::endl; - } + // if (!event.exists()) { + // std::cout << std::string(task.get_task_name()) << " " << task.index_point << " " << task.get_unique_id() << " " << task.orig_proc << " " << task.target_proc << " " << task.current_proc << " " << inflight.size() << std::endl; + // } // Assert that we found a valid event. taco_iassert(event.exists()); // Finally, trigger the event for anyone waiting on it. @@ -944,11 +944,11 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, const SelectMappingInput& input, SelectMappingOutput& output) { { - std::stringstream ss; - for (auto it : input.ready_tasks) { - ss << "{ " << it->get_task_name() << " " << it->get_unique_id() << "}, "; - } - std::cout << "Entering call to select_tasks_to_map, have tasks: " << ss.str() << std::endl; + // std::stringstream ss; + // for (auto it : input.ready_tasks) { + // ss << "{ " << it->get_task_name() << " " << it->get_unique_id() << "}, "; + // } + // std::cout << "Entering call to select_tasks_to_map, have tasks: " << ss.str() << std::endl; } if (!this->enableBackpressure) { DefaultMapper::select_tasks_to_map(ctx, input, output); @@ -988,13 +988,13 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, // launch loops go. // auto inflight = this->backPressureQueue[task->orig_proc]; auto inflight = this->backPressureQueue[task->target_proc]; - std::stringstream ss; - ss << "{ "; - for (auto it : this->backPressureQueue[task->target_proc]) { - ss << it.id << ", "; - } - ss << "}"; - std::cout << "Inflight for task: " << task->get_task_name() << " " << task->get_unique_id() << " " << ss.str() << std::endl; + // std::stringstream ss; + // ss << "{ "; + // for (auto it : this->backPressureQueue[task->target_proc]) { + // ss << it.id << ", "; + // } + // ss << "}"; + // std::cout << "Inflight for task: " << task->get_task_name() << " " << task->get_unique_id() << " " << ss.str() << std::endl; if (inflight.size() == this->maxInFlightTasks) { // We've hit the cap, so we can't schedule any more tasks. schedule = false; @@ -1009,32 +1009,32 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, } else { // Otherwise, we can schedule the task. Create a new event // and queue it up on the processor. - std::cout << "Adding event to execution queue " << std::string(task->get_task_name()) << " " << task->index_point << " " << task->orig_proc << " " << task->target_proc << " " << task->current_proc << " " << task->is_index_space << " " << task->index_point << std::endl; + // std::cout << "Adding event to execution queue " << std::string(task->get_task_name()) << " " << task->index_point << " " << task->orig_proc << " " << task->target_proc << " " << task->current_proc << " " << task->is_index_space << " " << task->index_point << std::endl; this->backPressureQueue[task->target_proc].push_back({ .id = task->get_unique_id(), .event = this->runtime->create_mapper_event(ctx), .schedTime = schedTime, }); - std::stringstream ss; - ss << "{ "; - for (auto it : this->backPressureQueue[task->target_proc]) { - ss << it.id << ", "; - } - ss << "}"; - std::cout << "target proc " << task->target_proc << " -> " << ss.str() << std::endl; + // std::stringstream ss; + // ss << "{ "; + // for (auto it : this->backPressureQueue[task->target_proc]) { + // ss << it.id << ", "; + // } + // ss << "}"; + // std::cout << "target proc " << task->target_proc << " -> " << ss.str() << std::endl; } } // Schedule tasks that are valid and have the target depth. if (schedule && (*it)->get_depth() == max_depth) { - std::cout << "Adding task to map task output: " << (*it)->get_task_name() << " " << (*it)->get_unique_id() << std::endl; - std::stringstream ss; - ss << "{"; - for (size_t idx = 0; idx < task->regions.size(); idx++) { - ss << Utilities::to_string(runtime, ctx, task->regions[idx], idx) << ", "; - } - ss << "}"; - std::cout << "Task has regs: " << ss.str() << std::endl; + // std::cout << "Adding task to map task output: " << (*it)->get_task_name() << " " << (*it)->get_unique_id() << std::endl; + // std::stringstream ss; + // ss << "{"; + // for (size_t idx = 0; idx < task->regions.size(); idx++) { + // ss << Utilities::to_string(runtime, ctx, task->regions[idx], idx) << ", "; + // } + // ss << "}"; + // std::cout << "Task has regs: " << ss.str() << std::endl; output.map_tasks.insert(*it); count++; } diff --git a/scripts/spbenchmark.py b/scripts/spbenchmark.py index e40809982..6303de173 100755 --- a/scripts/spbenchmark.py +++ b/scripts/spbenchmark.py @@ -57,7 +57,7 @@ def getCommand(self, tensor, benchKind, procs): BenchmarkKind.SpMV: ["-csr", self.getDISTALTensor(tensor, "csr")] + (["-tm:align128"] if self.gpu else []), BenchmarkKind.SpMSpV: ["-csc", self.getDISTALTensor(tensor, "csc"), "-spx", self.getUniformVec(tensor)], # TODO (rohany): Thread through the jdim here. - BenchmarkKind.SpMM: ["-tensor", self.getDISTALTensor(tensor, 'csr')] + (["-lg:eager_alloc_percentage", "5"] if self.gpu and tensor.name == "mycielskian19" else []), + BenchmarkKind.SpMM: ["-tensor", self.getDISTALTensor(tensor, 'csr')] + (["-lg:eager_alloc_percentage", "5"] if self.gpu and tensor.name == "mycielskian19" else []) + (["-tm:untrack_valid_regions", "-batched"] if self.gpu else []), # TODO (rohany): Thread through the jdim here. BenchmarkKind.SDDMM: ["-csr", self.getDISTALTensor(tensor, 'csr')], BenchmarkKind.SpAdd3: ["-tensorB", self.getDISTALTensor(tensor, 'csr'), From c87e0e0d81e8f932fee6926f8f58cb115fddca01 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Wed, 16 Feb 2022 10:02:12 -0800 Subject: [PATCH 08/12] fixup --- include/taco/index_notation/index_notation.h | 2 +- .../taco/index_notation/provenance_graph.h | 4 +-- legion/spmm/main.cpp | 7 ++-- legion/spmm/taco-generated.cu | 34 ++++++++++++------- legion/spmm/taco-generated.cuh | 6 ++-- src/index_notation/index_notation.cpp | 2 +- src/index_notation/provenance_graph.cpp | 17 +++++----- test/tests-distributed.cpp | 2 +- 8 files changed, 44 insertions(+), 30 deletions(-) diff --git a/include/taco/index_notation/index_notation.h b/include/taco/index_notation/index_notation.h index 2821abb40..34856bb96 100644 --- a/include/taco/index_notation/index_notation.h +++ b/include/taco/index_notation/index_notation.h @@ -570,7 +570,7 @@ class IndexStmt : public util::IntrusivePtr { /// variable, a \textit{tail strategy} is employed such as emitting a variable /// sized loop that handles remaining iterations. /// Preconditions: splitFactor is a positive nonzero integer - IndexStmt split(IndexVar i, IndexVar i1, IndexVar i2, size_t splitFactor) const; // TODO: TailStrategy + IndexStmt split(IndexVar i, IndexVar i1, IndexVar i2, ir::Expr splitFactor) const; // TODO: TailStrategy /// The divide transformation splits one index variable into /// two nested index variables, where the size of the outer diff --git a/include/taco/index_notation/provenance_graph.h b/include/taco/index_notation/provenance_graph.h index 91c53ed9c..e93909ccc 100644 --- a/include/taco/index_notation/provenance_graph.h +++ b/include/taco/index_notation/provenance_graph.h @@ -78,12 +78,12 @@ struct IndexVarRelNode : public util::Manageable, /// The split relation takes a parentVar's iteration space and stripmines into an outervar that iterates over splitFactor-sized /// iterations over innerVar struct SplitRelNode : public IndexVarRelNode { - SplitRelNode(IndexVar parentVar, IndexVar outerVar, IndexVar innerVar, size_t splitFactor); + SplitRelNode(IndexVar parentVar, IndexVar outerVar, IndexVar innerVar, ir::Expr splitFactor); const IndexVar& getParentVar() const; const IndexVar& getOuterVar() const; const IndexVar& getInnerVar() const; - const size_t& getSplitFactor() const; + const ir::Expr& getSplitFactor() const; void print(std::ostream& stream) const; bool equals(const SplitRelNode &rel) const; diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index c093b4854..9b1ba70ba 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -21,7 +21,7 @@ void top_level_task(const Task* task, const std::vector& regions bool dump = false, batched = 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, batchSize = 2; Realm::CommandLineParser parser; parser.add_option_string("-tensor", csrFileName); parser.add_option_bool("-dump", dump); @@ -30,6 +30,7 @@ void top_level_task(const Task* task, const std::vector& regions parser.add_option_int("-warmup", warmup); parser.add_option_int("-jdim", jDim); parser.add_option_bool("-batched", batched); + parser.add_option_int("-batchSize", batchSize); 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"; @@ -54,7 +55,7 @@ void top_level_task(const Task* task, const std::vector& regions partitionPackForcomputeLegionBatched packBatched; IndexPartition commPart; if (batched) { - packBatched = partitionForcomputeLegionBatched(ctx, runtime, &A, &B, &C, pieces); + packBatched = partitionForcomputeLegionBatched(ctx, runtime, &A, &B, &C, batchSize, pieces); commPart = createSparseAliasingPartitions(ctx, runtime, A.vals.get_index_space(), packBatched.APartition.valsPartition.get_index_partition()); } else { pack = partitionForcomputeLegion(ctx, runtime, &A, &B, &C, pieces); @@ -70,7 +71,7 @@ void top_level_task(const Task* task, const std::vector& regions auto avgTime = benchmarkAsyncCallWithWarmup(ctx, runtime, warmup, n, [&]() { if (dump) { runtime->fill_field(ctx, A.vals, A.valsParent, FID_VAL, valType(0)); } if (batched) { - computeLegionBatched(ctx, runtime, &A, &B, &C, &packBatched, pieces); + computeLegionBatched(ctx, runtime, &A, &B, &C, &packBatched, batchSize, pieces); 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); diff --git a/legion/spmm/taco-generated.cu b/legion/spmm/taco-generated.cu index a72c380a4..1a2be0932 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -25,6 +25,7 @@ struct task_2Args { int64_t B2Size; int64_t C2_dimension; int32_t gx; + int32_t gy; int64_t jo; int64_t pointID1; }; @@ -218,7 +219,7 @@ void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* } -partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx) { +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gy, int32_t gx) { RegionWrapper A_vals = A->vals; IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; size_t B1_dimension = B->dims[0]; @@ -236,7 +237,7 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + (gy - 1)) / gy); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -258,15 +259,15 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co B2CrdRect = B2CrdRect.make_empty(); } B2_crd_coloring[(*itr)] = B2CrdRect; - Point<2> AStart = Point<2>((0 / B2_dimension), (jo * 8)); - Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * 8 + 7),ADomain.hi()[1])); + Point<2> AStart = Point<2>((0 / B2_dimension), (jo * gy)); + Point<2> AEnd = Point<2>(TACO_MIN(((B1_dimension * B2_dimension - 1) / B2_dimension),ADomain.hi()[0]), TACO_MIN((jo * gy + (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<2> CStart = Point<2>(0, (jo * 8)); - Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * 8 + 7),CDomain.hi()[1])); + Point<2> CStart = Point<2>(0, (jo * gy)); + Point<2> CEnd = Point<2>(TACO_MIN((B1_dimension * B2_dimension - 1),CDomain.hi()[0]), TACO_MIN((jo * gy + (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(); @@ -313,13 +314,19 @@ partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Co computePartitions.CPartition.denseLevelRunPartitions = std::vector(2); computePartitions.CPartition.valsPartition = C_vals_partition; computePartitions.CPartition.denseLevelRunPartitions[0] = CDenseRun0Partition; + if (jo == 0) { + computePartitions.dummyReg = runtime->create_logical_region(ctx, fposoIndexSpace, LogicalRegion(A_vals).get_field_space()); + auto eqPart = runtime->create_equal_partition(ctx, computePartitions.dummyReg.get_index_space(), fposoIndexSpace); + computePartitions.dummyPart = runtime->get_logical_partition(ctx, computePartitions.dummyReg, eqPart); + runtime->fill_field(ctx, computePartitions.dummyReg, computePartitions.dummyReg, FID_VAL, double(0)); + } } return computePartitions; } __global__ -void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_blockStarts, int64_t pointID2, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorReduceNonExcldouble2 A_vals_red_accessor_non_excl, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t C2_dimension, int64_t jo, int64_t pointID1) { +void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_blockStarts, int64_t pointID2, AccessorRORect_1_1 B2_pos_accessor, AccessorROint32_t1 B2_crd_accessor, AccessorReduceNonExcldouble2 A_vals_red_accessor_non_excl, AccessorROdouble1 B_vals_ro_accessor, AccessorROdouble2 C_vals_ro_accessor, int64_t C2_dimension, int32_t gy, int64_t jo, int64_t pointID1) { int64_t block = blockIdx.x; int64_t thread = (threadIdx.x % (32)); @@ -355,8 +362,8 @@ void task_2DeviceKernel0(int64_t B2Size, int64_t fposo, int32_t gx, int64_t* i_b i_pos = i_pos + 1; i = i_pos; } - for (int64_t ji = 0; ji < 8; ji++) { - int64_t j = jo * 8 + ji; + for (int64_t ji = 0; ji < gy; ji++) { + int64_t j = jo * gy + ji; if (j >= C2_dimension) break; @@ -382,6 +389,7 @@ void task_2(const Task* task, const std::vector& regions, Contex int64_t B2Size = args->B2Size; int64_t C2_dimension = args->C2_dimension; int32_t gx = args->gx; + int32_t gy = args->gy; int64_t jo = args->jo; int64_t pointID1 = args->pointID1; @@ -410,11 +418,11 @@ void task_2(const Task* task, const std::vector& regions, Contex ); int64_t pointID2 = pointID1 * gx + fposo; if ((((B2Size + (gx - 1)) / gx + 2047) / 2048) > 0) { - task_2DeviceKernel0<<<(((B2Size + (gx - 1)) / gx + 2047) / 2048), (32 * 8)>>>(B2Size, fposo, gx, i_blockStarts, pointID2, B2_pos_accessor, B2_crd_accessor, A_vals_red_accessor_non_excl, B_vals_ro_accessor, C_vals_ro_accessor, C2_dimension, jo, pointID1); + task_2DeviceKernel0<<<(((B2Size + (gx - 1)) / gx + 2047) / 2048), (32 * 8)>>>(B2Size, fposo, gx, i_blockStarts, pointID2, B2_pos_accessor, B2_crd_accessor, A_vals_red_accessor_non_excl, B_vals_ro_accessor, C_vals_ro_accessor, C2_dimension, gy, jo, pointID1); } } -void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx) { +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gy, int32_t gx) { RegionWrapper A_vals = A->vals; auto A_vals_parent = A->valsParent; RegionWrapper B2_pos = B->indices[1][0]; @@ -429,7 +437,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; - for (int64_t jo = 0; jo < ((C2_dimension + 7) / 8); jo++) { + for (int64_t jo = 0; jo < ((C2_dimension + (gy - 1)) / gy); jo++) { int64_t pointID1 = jo + TACO_PARTITION_COLOR_OFFSET; Point<1> lowerBound = Point<1>(0); Point<1> upperBound = Point<1>((gx - 1)); @@ -439,6 +447,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT taskArgsRaw2.B2Size = B2Size; taskArgsRaw2.C2_dimension = C2_dimension; taskArgsRaw2.gx = gx; + taskArgsRaw2.gy = gy; taskArgsRaw2.jo = jo; taskArgsRaw2.pointID1 = pointID1; TaskArgument taskArgs = TaskArgument(&taskArgsRaw2, sizeof(task_2Args)); @@ -448,6 +457,7 @@ void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionT launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B2_crd), pointID1), 0, READ_ONLY, EXCLUSIVE, get_logical_region(B2_crd_parent)).add_field(FID_COORD)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(B_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, B_vals_parent).add_field(FID_VAL)); launcher.add_region_requirement(RegionRequirement(runtime->get_logical_partition_by_color(ctx, get_logical_region(C_vals), pointID1), 0, READ_ONLY, EXCLUSIVE, C_vals_parent).add_field(FID_VAL)); + launcher.add_region_requirement(RegionRequirement(partitionPack->dummyPart, 0, READ_WRITE, EXCLUSIVE, partitionPack->dummyReg).add_field(FID_VAL)); launcher.tag = launcher.tag | TACOMapper::UNTRACK_VALID_REGIONS; runtime->execute_index_space(ctx, launcher); diff --git a/legion/spmm/taco-generated.cuh b/legion/spmm/taco-generated.cuh index adf2ba615..c50215d41 100644 --- a/legion/spmm/taco-generated.cuh +++ b/legion/spmm/taco-generated.cuh @@ -12,6 +12,8 @@ struct partitionPackForcomputeLegionBatched { LegionTensorPartition APartition; LegionTensorPartition BPartition; LegionTensorPartition CPartition; + Legion::LogicalRegion dummyReg; + Legion::LogicalPartition dummyPart; }; @@ -20,9 +22,9 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg void computeLegion(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegion* partitionPack, int32_t gx); -partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gx); +partitionPackForcomputeLegionBatched partitionForcomputeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, int32_t gy, int32_t gx); -void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx); +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gy, int32_t gx); void registerTacoTasks(); #endif // TACO_GENERATED_CUH diff --git a/src/index_notation/index_notation.cpp b/src/index_notation/index_notation.cpp index 055156ea5..3376c4dff 100644 --- a/src/index_notation/index_notation.cpp +++ b/src/index_notation/index_notation.cpp @@ -1612,7 +1612,7 @@ IndexStmt IndexStmt::concretize() const { return stmt; } -IndexStmt IndexStmt::split(IndexVar i, IndexVar i1, IndexVar i2, size_t splitFactor) const { +IndexStmt IndexStmt::split(IndexVar i, IndexVar i1, IndexVar i2, ir::Expr splitFactor) const { IndexVarRel rel = IndexVarRel(new SplitRelNode(i, i1, i2, splitFactor)); string reason; diff --git a/src/index_notation/provenance_graph.cpp b/src/index_notation/provenance_graph.cpp index 63f4679f8..6f2339eed 100644 --- a/src/index_notation/provenance_graph.cpp +++ b/src/index_notation/provenance_graph.cpp @@ -138,10 +138,10 @@ struct SplitRelNode::Content { IndexVar parentVar; IndexVar outerVar; IndexVar innerVar; - size_t splitFactor; + ir::Expr splitFactor; }; -SplitRelNode::SplitRelNode(IndexVar parentVar, IndexVar outerVar, IndexVar innerVar, size_t splitFactor) +SplitRelNode::SplitRelNode(IndexVar parentVar, IndexVar outerVar, IndexVar innerVar, ir::Expr splitFactor) : IndexVarRelNode(SPLIT), content(new Content) { content->parentVar = parentVar; content->outerVar = outerVar; @@ -158,7 +158,7 @@ const IndexVar& SplitRelNode::getOuterVar() const { const IndexVar& SplitRelNode::getInnerVar() const { return content->innerVar; } -const size_t& SplitRelNode::getSplitFactor() const { +const ir::Expr& SplitRelNode::getSplitFactor() const { return content->splitFactor; } @@ -193,7 +193,7 @@ std::vector SplitRelNode::computeRelativeBound(std::set defi return parentBound; // splitting pos space does not change coordinate bounds } - ir::Expr splitFactorLiteral = ir::Literal::make(getSplitFactor(), variableExprs[getParentVar()].type()); + ir::Expr splitFactorLiteral = getSplitFactor(); if (!outerVarDefined && !innerVarDefined) { return parentBound; @@ -230,6 +230,7 @@ std::vector SplitRelNode::deriveIterBounds(taco::IndexVar indexVar, std::vector parentBound = parentIterBounds.at(getParentVar()); Datatype splitFactorType = parentBound[0].type(); + auto splitFactor = this->getSplitFactor(); if (indexVar == getOuterVar()) { // The outer variable must always range from 0 to the extent of the bounds (chunked up). // This is a noop for the common case where all of our loops start at 0. However, it is @@ -237,11 +238,11 @@ std::vector SplitRelNode::deriveIterBounds(taco::IndexVar indexVar, // and instead start at a position like T_pos[i].lo. ir::Expr minBound = 0; auto upper = ir::Sub::make(parentBound[1], parentBound[0]); - ir::Expr maxBound = ir::Div::make(ir::Add::make(upper, ir::Literal::make(getSplitFactor()-1, splitFactorType)), ir::Literal::make(getSplitFactor(), splitFactorType)); + ir::Expr maxBound = ir::Div::make(ir::Add::make(upper, ir::Sub::make(splitFactor, 1)), splitFactor); return {minBound, maxBound}; } else if (indexVar == getInnerVar()) { ir::Expr minBound = 0; - ir::Expr maxBound = ir::Literal::make(getSplitFactor(), splitFactorType); + ir::Expr maxBound = this->getSplitFactor(); return {minBound, maxBound}; } taco_ierror; @@ -257,7 +258,7 @@ ir::Expr SplitRelNode::recoverVariable(taco::IndexVar indexVar, // Include the lower bound of the variable being recovered. Normally, this is 0, but // for cases like DIVIDE_ONTO_PARTITION, it is not. return ir::Add::make( - ir::Add::make(ir::Mul::make(variableNames[getOuterVar()], ir::Literal::make(getSplitFactor(), splitFactorType)), variableNames[getInnerVar()]), + ir::Add::make(ir::Mul::make(variableNames[getOuterVar()], this->getSplitFactor()), variableNames[getInnerVar()]), parentIterBounds[indexVar][0] ); } @@ -282,7 +283,7 @@ ir::Stmt SplitRelNode::recoverChild(IndexVar indexVar, std::vector def else { // innerVar = parentVar - outerVar * splitFactor ir::Expr subStmt = ir::Sub::make(variableNames[getParentVar()], - ir::Mul::make(variableNames[getOuterVar()], ir::Literal::make(getSplitFactor(), splitFactorType))); + ir::Mul::make(variableNames[getOuterVar()], this->getSplitFactor())); if (emitVarDecl) { return ir::Stmt(ir::VarDecl::make(variableNames[getInnerVar()], subStmt)); } diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp index 4ae97dd6e..491b9df66 100644 --- a/test/tests-distributed.cpp +++ b/test/tests-distributed.cpp @@ -1864,7 +1864,7 @@ TEST(distributed, legionSpMM) { const int BATCHED_NNZ_PER_WARP = BATCHED_NNZ_PER_THREAD * WARP_SIZE; const int BATCHED_NNZ_PER_TB = BATCHED_NNZ_PER_THREAD * BLOCK_SIZE; auto cpuConsLowered = lowerLegionSeparatePartitionCompute(cpuCons, "computeLegionBatched", false /* waitOnFutureMap */); - auto gpuCons = stmt.split(j, jo, ji, BATCH_SIZE) + auto gpuCons = stmt.split(j, jo, ji, gy) .reorder({jo, i, k, ji}) .fuse(i, k, f) .pos(f, fpos, B(i, k)) From 48fb91ccddc30649ee382ca88fab7e9e3cd8a5f2 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Wed, 16 Feb 2022 11:42:35 -0800 Subject: [PATCH 09/12] fixup --- legion/src/taco_mapper.cpp | 3 ++- scripts/spbenchmark.py | 27 ++++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/legion/src/taco_mapper.cpp b/legion/src/taco_mapper.cpp index d9150a594..4cb3b1387 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -864,7 +864,8 @@ void TACOMapper::slice_task(const Legion::Mapping::MapperContext ctx, // std::cout << it << " " << std::endl; // } // std::cout << std::endl; - bool recurse = ((task.tag & BACKPRESSURE_TASK) != 0) && input.domain.get_volume() > 1; + // bool recurse = ((task.tag & BACKPRESSURE_TASK) != 0) && input.domain.get_volume() > 1; + bool recurse = false; switch (input.domain.get_dim()) { #define BLOCK(DIM) \ case DIM: \ diff --git a/scripts/spbenchmark.py b/scripts/spbenchmark.py index 6303de173..176743d75 100755 --- a/scripts/spbenchmark.py +++ b/scripts/spbenchmark.py @@ -103,9 +103,34 @@ def getCommand(self, tensor, benchKind, procs): "-ll:fsize", fbSize, "-pieces", str(procs), ] + + def getSpMMBatchSize(): + if tensor.name in ["arabic-2005", "mycielskian19", "nlpkkt240"]: + return 32 + elif tensor.name == "it-2004": + return 16 + elif tensor.name == "kmer_A2a": + assert(False) + elif tensor.name == "kmer_V1r": + assert(False) + elif tensor.name == "sk-2005": + # TODO (rohany): Play around with batch size at diff node counts. + return 4 + elif tensor.name == "twitter7": + # TODO (rohany): Play around with batch size at diff node counts. + return 4 + elif tensor.name == "uk-2005": + return 16 + elif tensor.name == "webbase-2001": + assert(False) + else: + assert(False) + spmmBatchArgs = [] + if benchKind == BenchmarkKind.SpMM: + spmmBatchArgs = ["-batchSize", str(getSpMMBatchSize())] assert(benchKind in args) - return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + spmmBatchArgs def getBinary(self, benchKind): tacoDir = os.environ.get("TACO_BUILD_DIR", None) From 93d6863f55b980df14f8578e78dde0d4c0246467 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Thu, 17 Feb 2022 09:47:35 -0800 Subject: [PATCH 10/12] fixup --- legion/legion | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/legion/legion b/legion/legion index a72f18ba2..0549d91c9 160000 --- a/legion/legion +++ b/legion/legion @@ -1 +1 @@ -Subproject commit a72f18ba27241054e4930161ed4dde181c05ad46 +Subproject commit 0549d91c944ff717c430a48b1a2df32672d5460a From 117e42d7249357926b52d4e2744926490da2d941 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Tue, 22 Feb 2022 11:00:18 -0800 Subject: [PATCH 11/12] fixup --- legion/spmm/main.cpp | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 9b1ba70ba..4d6f693ed 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -18,7 +18,7 @@ typedef double valType; void top_level_task(const Task* task, const std::vector& regions, Context ctx, Runtime* runtime) { std::string csrFileName; - bool dump = false, batched = false; + bool dump = false, batched = false, CCpu = 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, batchSize = 2; @@ -31,6 +31,7 @@ void top_level_task(const Task* task, const std::vector& regions parser.add_option_int("-jdim", jDim); parser.add_option_bool("-batched", batched); parser.add_option_int("-batchSize", batchSize); + parser.add_option_bool("-Ccpu", CCpu); 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"; @@ -65,7 +66,26 @@ void top_level_task(const Task* task, const std::vector& regions // Create a long lasting instance for the C region to force communcation costs. if (batched) { - tacoFill(ctx, runtime, C.vals, valType(1)); + auto colorSpace = runtime->create_index_space(ctx, Rect<1>(0, pieces - 1)); + taco_iassert(pieces <= 32); + auto dom = runtime->get_index_space_domain(ctx, C.vals.get_index_space()); + auto size = jDim / pieces; + DomainPointColoring coloring; + for (int i = 0; i < pieces; i++) { + coloring[i] = Rect<2>({dom.lo()[0], i * size}, {dom.hi()[0], (i+1) * size - 1}); + } + auto ipart = runtime->create_index_partition(ctx, C.valsParent.get_index_space(), Rect<1>(0, pieces - 1), coloring); + auto lpart = runtime->get_logical_partition(ctx, C.vals, ipart); + tacoFill(ctx, runtime, C.vals, lpart, valType(1)); + if (CCpu) { + auto numOMPs = runtime->select_tunable_value(ctx, Mapping::DefaultMapper::DEFAULT_TUNABLE_GLOBAL_OMPS).get(); + // TODO (rohany): I'm hacking here, but this distribution could be easily done using + // DISTAL's data distribution language. + IndexTaskLauncher launcher(TID_DUMMY_READ_REGION, Rect<1>(0, numOMPs - 1), TaskArgument(), ArgumentMap()); + launcher.add_region_requirement(RegionRequirement(C.vals, READ_ONLY, EXCLUSIVE, C.valsParent).add_field(FID_VAL)); + launcher.tag |= TACOMapper::MAP_TO_OMP_OR_LOC; + runtime->execute_index_space(ctx, launcher).wait_all_results(); + } } auto avgTime = benchmarkAsyncCallWithWarmup(ctx, runtime, warmup, n, [&]() { From a60d622931c6b01f5bf105079e7d6901146649ac Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Tue, 22 Feb 2022 11:00:59 -0800 Subject: [PATCH 12/12] fixup --- scripts/spbenchmark.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/scripts/spbenchmark.py b/scripts/spbenchmark.py index 176743d75..279bfc3c3 100755 --- a/scripts/spbenchmark.py +++ b/scripts/spbenchmark.py @@ -108,11 +108,14 @@ def getSpMMBatchSize(): if tensor.name in ["arabic-2005", "mycielskian19", "nlpkkt240"]: return 32 elif tensor.name == "it-2004": - return 16 + if procs <= 4: + return 16 + else: + return 32 elif tensor.name == "kmer_A2a": - assert(False) + return 4 elif tensor.name == "kmer_V1r": - assert(False) + return 4 elif tensor.name == "sk-2005": # TODO (rohany): Play around with batch size at diff node counts. return 4 @@ -122,12 +125,14 @@ def getSpMMBatchSize(): elif tensor.name == "uk-2005": return 16 elif tensor.name == "webbase-2001": - assert(False) + return 4 else: assert(False) spmmBatchArgs = [] if benchKind == BenchmarkKind.SpMM: - spmmBatchArgs = ["-batchSize", str(getSpMMBatchSize())] + spmmBatchArgs = ["-batchSize", str(getSpMMBatchSize())]#, "-tm:fill_cpu"] + if 'SPMM_CPU_C' in os.environ: + spmmBatchArgs += ["-tm:fill_cpu", "-Ccpu"] assert(benchKind in args) return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + spmmBatchArgs