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/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..5949eca1d 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 { @@ -201,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 @@ -209,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/legion b/legion/legion index a72f18ba2..0549d91c9 160000 --- a/legion/legion +++ b/legion/legion @@ -1 +1 @@ -Subproject commit a72f18ba27241054e4930161ed4dde181c05ad46 +Subproject commit 0549d91c944ff717c430a48b1a2df32672d5460a diff --git a/legion/spmm/main.cpp b/legion/spmm/main.cpp index 75a954004..4d6f693ed 100644 --- a/legion/spmm/main.cpp +++ b/legion/spmm/main.cpp @@ -6,17 +6,22 @@ #include "legion_utils.h" #include "legion_string_utils.h" #include "error.h" + +#ifdef TACO_USE_CUDA +#include "taco-generated.cuh" +#else #include "taco-generated.h" +#endif using namespace Legion; typedef double valType; void top_level_task(const Task* task, const std::vector& regions, Context ctx, Runtime* runtime) { std::string csrFileName; - bool dump = false; + bool dump = false, 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; + 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); @@ -24,6 +29,9 @@ 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); + 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"; @@ -41,14 +49,53 @@ 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, 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); + 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); + // Create a long lasting instance for the C region to force communcation costs. + if (batched) { + 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, [&]() { 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, 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); + } #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 @@ -81,6 +128,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/spmm/taco-generated.cpp b/legion/spmm/taco-generated.cpp index 996552872..dde4eca29 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; @@ -18,19 +19,36 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t B2Size; + int32_t gx; + int64_t pointID1; +}; + +struct task_3Args { + int64_t A2_dimension; + 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; 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(); + int64_t B2Size = runtime->get_index_space_domain(ctx, get_index_space(B2_crd)).hi()[0] + 1; Point<1> lowerBound = Point<1>(0); @@ -60,13 +78,12 @@ 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)); - auto computePartitions = partitionPackForcomputeLegion(); computePartitions.APartition.indicesPartitions = std::vector>(2); computePartitions.APartition.denseLevelRunPartitions = std::vector(2); computePartitions.APartition.valsPartition = A_vals_partition; @@ -106,7 +123,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; @@ -133,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; @@ -162,14 +179,300 @@ 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); } + +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]; + 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]; + 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]; + 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]; + + 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 + 7) / 8); 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 AColoring = DomainPointColoring(); + 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> 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 * 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(); + } + 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).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] = 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; +} + +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]; + 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_3Args* args = (task_3Args*)(task->args); + int64_t A2_dimension = args->A2_dimension; + 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 = 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)); + 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; + for (int64_t ji = 0; ji < 8; ji++) { + int64_t j = jo * 8 + ji; + if (j >= C2_dimension) + continue; + + 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)]; + } + } +} + +void computeLegionBatched(Legion::Context ctx, Legion::Runtime* runtime, LegionTensor* A, LegionTensor* B, LegionTensor* C, partitionPackForcomputeLegionBatched* partitionPack, int32_t gx) { + 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]; + 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; + 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 + 7) / 8); 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_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(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); + + } +} 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"); + } + { + 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 3b805b815..1a2be0932 100644 --- a/legion/spmm/taco-generated.cu +++ b/legion/spmm/taco-generated.cu @@ -21,6 +21,15 @@ struct task_1Args { int32_t gx; }; +struct task_2Args { + int64_t B2Size; + int64_t C2_dimension; + int32_t gx; + int32_t gy; + 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 +39,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 +51,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 +72,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 +81,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 +96,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 +105,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 +181,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); } } @@ -177,7 +193,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; @@ -198,9 +214,255 @@ 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 gy, int32_t gx) { + RegionWrapper A_vals = A->vals; + IndexSpace A_dense_run_0 = A->denseLevelRuns[0]; + 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]; + size_t 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 + (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)); + 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 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> 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 * 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(); + } + 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).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); + 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; + 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, int32_t gy, 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 < gy; ji++) { + int64_t j = jo * gy + 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; + int32_t gy = args->gy; + 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, gy, jo, pointID1); + } +} + +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]; + 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; + 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 + (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)); + 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.gy = gy; + 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.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); + + } +} void registerTacoTasks() { { TaskVariantRegistrar registrar(taskID(1), "task_1"); @@ -208,4 +470,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..c50215d41 100644 --- a/legion/spmm/taco-generated.cuh +++ b/legion/spmm/taco-generated.cuh @@ -8,10 +8,23 @@ struct partitionPackForcomputeLegion { LegionTensorPartition BPartition; }; +struct partitionPackForcomputeLegionBatched { + LegionTensorPartition APartition; + LegionTensorPartition BPartition; + LegionTensorPartition CPartition; + Legion::LogicalRegion dummyReg; + Legion::LogicalPartition dummyPart; +}; + 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 gy, 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/legion/spmm/taco-generated.h b/legion/spmm/taco-generated.h index 6b952b5f0..ad2fd96f6 100644 --- a/legion/spmm/taco-generated.h +++ b/legion/spmm/taco-generated.h @@ -8,10 +8,22 @@ 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_H 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..4cb3b1387 100644 --- a/legion/src/taco_mapper.cpp +++ b/legion/src/taco_mapper.cpp @@ -15,17 +15,21 @@ 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(); + // std::cout << "ONLY 1 MAPPER PER NODE" << std::endl; #ifdef TACO_USE_LOGGING_MAPPER runtime->replace_default_mapper(new Mapping::LoggingWrapper(new TACOMapper(runtime->get_mapper_runtime(), machine, proc, TACOMapperName)), Processor::NO_PROC); #else @@ -139,25 +143,301 @@ 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, + 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); + // 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. 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) { + // 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); } } } // 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(); } } @@ -579,6 +859,13 @@ 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; + bool recurse = false; switch (input.domain.get_dim()) { #define BLOCK(DIM) \ case DIM: \ @@ -587,7 +874,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; \ } @@ -597,11 +884,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. @@ -613,7 +912,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++) { @@ -623,6 +930,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. @@ -634,6 +944,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 { @@ -663,11 +980,22 @@ void TACOMapper::select_tasks_to_map(const MapperContext ctx, { auto task = *it; bool schedule = true; - if ((task->tag & BACKPRESSURE_TASK) != 0) { + if (std::string(task->get_task_name()) == "task_2") { + taco_iassert((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. - 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; @@ -682,16 +1010,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++; } diff --git a/scripts/spbenchmark.py b/scripts/spbenchmark.py index e40809982..279bfc3c3 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'), @@ -103,9 +103,39 @@ 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": + if procs <= 4: + return 16 + else: + return 32 + elif tensor.name == "kmer_A2a": + return 4 + elif tensor.name == "kmer_V1r": + return 4 + 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": + return 4 + else: + assert(False) + spmmBatchArgs = [] + if benchKind == BenchmarkKind.SpMM: + 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] + return lassenPrefix + [self.getBinary(benchKind)] + legionArgs + commonArgs + args[benchKind] + spmmBatchArgs def getBinary(self, benchKind): tacoDir = os.environ.get("TACO_BUILD_DIR", None) 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/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/src/lower/lowerer_impl.cpp b/src/lower/lowerer_impl.cpp index b9c8eccdd..f787b9cb8 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,56 @@ 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::map> partitionedByOuter; + 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 (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) { + 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. I think I know what to do here now? + // taco_iassert(!(isPartitionedByOuterVars && affineProjection.defined())); + } auto coloringLoop = ir::For::make( domainIter, ir::Call::make(pointInDimT.getName(), {domain}, pointInDimT), @@ -2190,6 +2240,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 +2349,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; } @@ -2366,29 +2429,44 @@ 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); - // 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) && !affineProjection.defined()) { + 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 sparseGatherProjection = this->constructSparseGatherProjection(posAccess, tvAccess, posMode); + if (affineProjection.defined()) { + 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. + 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 +2714,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 +5787,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)) { @@ -6017,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 @@ -6308,6 +6389,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 +6748,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 +7064,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..491b9df66 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,43 @@ 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 = 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. + .reorder({jo, i, k, ji}) + .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; + 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, gy) + .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) {