*: support a batched schedule for SpMM that uses less memory

include/taco/index_notation/index_notation.h

@@ -570,7 +570,7 @@ class IndexStmt : public util::IntrusivePtr<const IndexStmtNode> {
   /// 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

include/taco/index_notation/provenance_graph.h

@@ -78,12 +78,12 @@ struct IndexVarRelNode : public util::Manageable<IndexVarRelNode>,
 /// 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;

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

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,19 +202,29 @@ 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
   // dimensional space is projected onto a higher dimensional space. For dimensions that
   // 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<int> overs) {
+    for (auto it : overs) {
+      this->overrides.push_back(it);
+    }
+    return *this;
+  }
+  template <typename ... Overs>
+  AffineProjection addOverrides(const Overs&... overs) { return addOverrides(std::vector<int>{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<int> projs;
+  std::vector<int> overrides;
 };
 
 // Helper class to perform partitions for the downwards partition path for

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<std::tuple<Legion::Memory::Kind, Legion::IndexSpace, Legion::FieldSpace, Legion::ReductionOpID>, Legion::LayoutConstraintID> distalReductionConstraintCache;
 
+  std::map<std::tuple<unsigned long long, size_t, Legion::Memory>, std::vector<Legion::Mapping::PhysicalInstance>> 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:

legion/spmm/main.cpp

@@ -6,24 +6,32 @@
 #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<PhysicalRegion>& 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);
   parser.add_option_int("-n", n);
   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<PhysicalRegion>& 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<size_t>();
+      // 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<valType>();
   Runtime::add_registration_callback(register_taco_mapper);
   Runtime::preregister_sharding_functor(TACOShardingFunctorID, new TACOShardingFunctor());
   return Runtime::start(argc, argv);