Rematerialization via instances attribute.

Support generating 0 or multiple instances of a single operation in the source file by replacing the `decisions` attribute by an `instances` attribute. The `instances` attribute is an array where each entry represents a different instance of the operation in the generated code.

For now, only the first instance is actually checked and generated. Lowering passes expect a single instance to be present and will emit an error otherwise. A new pass ensures that all ComputeOp have at least one instance when assigning default attribute.

PiperOrigin-RevId: 395042740

Author: ulysseB

canonicalization_patterns.cc

@@ -132,15 +132,16 @@ class SimplifySairOperands : public RewritePattern {
 // Remove duplicate inputs and duplicate outputs of sair.map operations.
 mlir::LogicalResult DeduplicateMapInputsOutputs(
     SairMapOp op, mlir::PatternRewriter &rewriter) {
+  if (op.HasCopies()) return mlir::failure();
+
   int domain_size = op.domain().size();
   llvm::SmallVector<mlir::Value> new_operands;
   llvm::SmallVector<mlir::Attribute> new_mappings;
 
   llvm::SmallVector<mlir::Value> old_results_to_keep;
   llvm::SmallVector<mlir::Value> new_scalar_results;
   llvm::SmallVector<mlir::Type> new_result_types;
-  llvm::SmallVector<mlir::Attribute> new_storages;
-  llvm::SmallVector<mlir::Attribute> new_copies;
+  llvm::SmallBitVector remaining_outputs(op->getNumResults());
 
   std::vector<int> block_args_to_erase;
   for (ValueOperand operand : op.ValueOperands()) {
@@ -181,7 +182,15 @@ mlir::LogicalResult DeduplicateMapInputsOutputs(
     // Deduplicate results.
     for (int j = 0; j < i; ++j) {
       if (scalar_value != return_op.getOperand(j)) continue;
-      if (op.Storage(i) != op.Storage(j)) continue;
+      bool same_storage = true;
+      for (int k = 0, e = op.NumInstances(); k < e; ++k) {
+        if (op.GetDecisions(k).storage() != op.GetDecisions(j).storage()) {
+          same_storage = false;
+          break;
+        }
+      }
+      if (!same_storage) continue;
+
       // Don't deduplicate with dead results that will be removed.
       if (op.getResult(j).use_empty()) continue;
       result.replaceAllUsesWith(op.getResult(j));
@@ -196,10 +205,7 @@ mlir::LogicalResult DeduplicateMapInputsOutputs(
     old_results_to_keep.push_back(result);
     new_scalar_results.push_back(scalar_value);
     new_result_types.push_back(result.getType());
-    mlir::Attribute new_storage = op.Storage(i);
-    new_storages.push_back(new_storage == nullptr ? rewriter.getUnitAttr()
-                                                  : new_storage);
-    new_copies.push_back(rewriter.getArrayAttr(op.GetCopies(i)));
+    remaining_outputs.set(i);
   }
 
   // Create the new operation if necessary.
@@ -213,15 +219,14 @@ mlir::LogicalResult DeduplicateMapInputsOutputs(
   rewriter.eraseOp(return_op);
 
   rewriter.setInsertionPoint(op);
-  DecisionsAttr decisions = op.GetDecisions();
-  auto new_decisions =
-      DecisionsAttr::get(decisions.sequence(), decisions.loop_nest(),
-                         rewriter.getArrayAttr(new_storages),
-                         decisions.expansion(), op.getContext());
+  mlir::ArrayAttr new_instances = MkArrayAttrMapper<DecisionsAttr>(
+      MapStorage(MkArrayAttrFilter(remaining_outputs)))(op.instancesAttr());
+  mlir::ArrayAttr new_copies =
+      MkArrayAttrFilter(remaining_outputs)(op.copiesAttr());
   SairMapOp new_op = rewriter.create<SairMapOp>(
       op.getLoc(), new_result_types, op.domain(),
       rewriter.getArrayAttr(new_mappings), new_operands, op.shape(),
-      new_decisions, rewriter.getArrayAttr(new_copies));
+      new_instances, new_copies);
   new_op.body().takeBody(op.body());
 
   for (auto [old_res, new_res] :

expansion.cc

@@ -90,7 +90,7 @@ class CopyExpansionPattern : public ExpansionPattern {
 
 mlir::LogicalResult CopyExpansionPattern::Match(ComputeOpInstance op) const {
   if (op.is_copy()) return mlir::success();
-  ComputeOp compute_op = op.AsComputeOp();
+  ComputeOp compute_op = op.GetComputeOp();
   return mlir::success(isa<SairCopyOp>(compute_op.getOperation()));
 }
 

expansion.h

@@ -66,7 +66,7 @@ class TypedExpansionPattern : public ExpansionPattern {
 
   mlir::LogicalResult Match(ComputeOpInstance op) const final {
     if (op.is_copy()) return mlir::failure();
-    auto cast_op = dyn_cast<OpTy>(*op.AsComputeOp());
+    auto cast_op = dyn_cast<OpTy>(*op.GetComputeOp());
     if (cast_op == nullptr) return mlir::failure();
     return Match(cast_op);
   }

sair_attributes.cc

@@ -1337,6 +1337,30 @@ mlir::LogicalResult VerifyMappingShape(const AttrLocation &loc,
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// DecisionsAttr
+//===----------------------------------------------------------------------===//
+
+std::function<DecisionsAttr(DecisionsAttr)> MapLoopNest(
+    std::function<mlir::ArrayAttr(mlir::ArrayAttr)> loop_nest_fn) {
+  return [loop_nest_fn](DecisionsAttr decisions) -> DecisionsAttr {
+    if (decisions == nullptr) return nullptr;
+    return DecisionsAttr::get(
+        decisions.sequence(), loop_nest_fn(decisions.loop_nest()),
+        decisions.storage(), decisions.expansion(), decisions.getContext());
+  };
+}
+
+std::function<DecisionsAttr(DecisionsAttr)> MapStorage(
+    std::function<mlir::ArrayAttr(mlir::ArrayAttr)> storage_fn) {
+  return [storage_fn](DecisionsAttr decisions) -> DecisionsAttr {
+    if (decisions == nullptr) return nullptr;
+    return DecisionsAttr::get(decisions.sequence(), decisions.loop_nest(),
+                              storage_fn(decisions.storage()),
+                              decisions.expansion(), decisions.getContext());
+  };
+}
+
 }  // namespace sair
 
 #include "sair_structs.cc.inc"

sair_attributes.h

@@ -556,4 +556,23 @@ class MappingUnStripeExpr
 using namespace mlir;  // NOLINT
 #include "sair_structs.h.inc"
 
+namespace sair {
+
+// Below are helper functions to manipulate DecisionsAttr. Each helper takes a
+// function and returns a function that applies the first function to a field of
+// a DecisionsAttr. Helpers return functions rather than directly applying the
+// transformation so that it is easier to combine transformations.
+
+// Takes a function that updates a loop nest and returns a function that updates
+// the loop nest field of a DecisionsAttr.
+std::function<DecisionsAttr(DecisionsAttr)> MapLoopNest(
+    std::function<mlir::ArrayAttr(mlir::ArrayAttr)> loop_nest_fn);
+
+// Takes a function that updates a list of storages and returns a function that
+// updates the storage field of a DecisionsAttr.
+std::function<DecisionsAttr(DecisionsAttr)> MapStorage(
+    std::function<mlir::ArrayAttr(mlir::ArrayAttr)> storage_fn);
+
+}  // namespace sair
+
 #endif  // SAIR_SAIR_ATTRIBUTES_H_

sair_base.td

@@ -215,11 +215,13 @@ def SairDecisionsAttr : SairStructAttr<"DecisionsAttr", [
   StructFieldAttr<"expansion", OptionalAttr<StrAttr>>
 ]>;
 
-// An attribute that species copies of the result of an operation.
+// An attribute that specifies instances of an operation.
+def SairInstancesAttr
+  : TypedArrayAttrBase<SairDecisionsAttr, "array of Sair decisions">;
+
+// An attribute that specifies copies of the result of an operation.
 def SairCopiesAttr
-  : TypedArrayAttrBase<
-      TypedArrayAttrBase<SairDecisionsAttr, "array of Sair decisions">,
-      "array of array of Sair decisions">;
+  : TypedArrayAttrBase<SairInstancesAttr, "array of array of Sair decisions">;
 
 //===----------------------------------------------------------------------===//
 // Sair Types
@@ -352,8 +354,7 @@ def SairOpInterface : OpInterface<"SairOp"> {
         dependencies accross iterations for the given !sair.value operand}],
       "llvm::SmallBitVector", "CarryingDimensions",
       (ins "int":$sair_operand), [{}], [{
-        int size = cast<ConcreteOp>(this->getOperation()).domain().size();
-        return llvm::SmallBitVector(size);
+        return llvm::SmallBitVector($_op.domain().size());
       }]
     >,
     InterfaceMethod<
@@ -370,7 +371,12 @@ def SairOpInterface : OpInterface<"SairOp"> {
       [{Indicates the size of each sub-domain. The first sub-domain is always
         the parallel domain}],
       "llvm::SmallVector<int, 2>", "SubDomains", (ins), [{}]
-    >
+    >,
+    InterfaceMethod<
+      [{Indicates if the operation has exactly one instance and no copy.}],
+      "bool", "HasExactlyOneInstance", (ins), [{}], [{
+        return sair::HasExactlyOneInstance(this->getOperation());
+      }]>
   ];
 
   let verify = [{return VerifySairOp(op);}];
@@ -410,34 +416,59 @@ def SairComputeOp : OpInterface<"ComputeOp"> {
 
   let methods = [
     InterfaceMethod<
-      "Returns lowering decisions for the operation",
-      "DecisionsAttr", "GetDecisions", (ins), [{}], [{
-        if (!$_op.decisions().hasValue()) {
-          return DecisionsAttr::get(
-            nullptr, nullptr, nullptr, nullptr, $_op.getContext());
-        }
-        return *$_op.decisions();
+      "Returns the number of instances of the operation",
+      "int", "NumInstances", (ins), [{}], [{
+        llvm::Optional<mlir::ArrayAttr> instances = $_op.instances();
+        if (!instances.hasValue()) return 0;
+        return instances.getValue().size();
+      }]>,
+    InterfaceMethod<
+      "Returns lowering decisions for the given operation instance",
+      "DecisionsAttr", "GetDecisions", (ins "int":$instance), [{}], [{
+        mlir::ArrayAttr instances = $_op.instances().getValue();
+        return instances.getValue()[instance].cast<DecisionsAttr>();
       }]
     >,
     InterfaceMethod<
       "Sets lowering decisions for the operation",
-      "void", "SetDecisions", (ins "DecisionsAttr":$value), [{}], [{
-        $_op.decisionsAttr(value);
+      "void", "SetDecisions", (ins "int":$instance, "DecisionsAttr":$value),
+      [{}], [{
+        auto instances = llvm::to_vector<4>(
+          $_op.instances().getValue().getValue());
+        instances[instance] = value;
+        $_op.instancesAttr(mlir::ArrayAttr::get($_op.getContext(), instances));
       }]
     >,
+    InterfaceMethod<
+      "Appends an instance to the list of instances",
+      "void", "AddInstance", (ins "DecisionsAttr":$value),
+      [{}], [{
+        llvm::SmallVector<mlir::Attribute> instances;
+        if ($_op.instances().hasValue()) {
+          instances = llvm::to_vector<4>(
+            $_op.instances().getValue().getValue());
+        }
+        instances.push_back(value);
+        $_op.instancesAttr(mlir::ArrayAttr::get($_op.getContext(), instances));
+      }]
+    >,
+
+    // TODO(ulysse): Legacy interface that returns decisions for the first
+    // instance. This will be removed.
     InterfaceMethod<
       "Returns the loop nest to generate when lowering the operation",
       "llvm::Optional<mlir::ArrayAttr>", "loop_nest", (ins), [{}], [{
-        if (!$_op.decisions().hasValue()) return llvm::None;
-        if ($_op.decisions()->loop_nest() == nullptr) return llvm::None;
-        return $_op.decisions()->loop_nest();
+        if ($_op.NumInstances() < 1) return llvm::None;
+        DecisionsAttr decisions = $_op.GetDecisions(0);
+        if (decisions.loop_nest() == nullptr) return llvm::None;
+        return decisions.loop_nest();
       }]
     >,
     InterfaceMethod<
       "Sets the loop nest to generate when lowering the operation",
       "void", "setLoopNest", (ins "mlir::ArrayAttr":$loop_nest), [{
-        DecisionsAttr decisions = $_op.GetDecisions();
-        $_op.SetDecisions(DecisionsAttr::get(
+        DecisionsAttr decisions = $_op.GetDecisions(0);
+        $_op.SetDecisions(0, DecisionsAttr::get(
           decisions.sequence(), loop_nest, decisions.storage(),
           decisions.expansion(), $_op.getContext()));
       }]
@@ -455,9 +486,10 @@ def SairComputeOp : OpInterface<"ComputeOp"> {
     InterfaceMethod<
       "Returns the storage of the values produced by the operation",
       "Optional<mlir::ArrayAttr>", "storage", (ins), [{}], [{
-        if (!$_op.decisions().hasValue()) return llvm::None;
-        if ($_op.decisions()->storage() == nullptr) return llvm::None;
-        return $_op.decisions()->storage();
+        if ($_op.NumInstances() < 1) return llvm::None;
+        DecisionsAttr decisions = $_op.GetDecisions(0);
+        if (decisions.storage() == nullptr) return llvm::None;
+        return decisions.storage();
       }]
     >,
     InterfaceMethod<
@@ -472,7 +504,7 @@ def SairComputeOp : OpInterface<"ComputeOp"> {
     InterfaceMethod<
       "Sets the storage of the values produced by the operation",
       "void", "SetStorage", (ins "int":$result, "BufferAttr":$buffer), [{}], [{
-        DecisionsAttr decisions = $_op.GetDecisions();
+        DecisionsAttr decisions = $_op.GetDecisions(0);
         llvm::SmallVector<mlir::Attribute> values;
         int num_results = $_op->getNumResults();
         if (decisions.storage() == nullptr) {
@@ -482,25 +514,26 @@ def SairComputeOp : OpInterface<"ComputeOp"> {
         }
         values[result] = buffer;
         auto new_storage = mlir::ArrayAttr::get($_op.getContext(), values);
-        $_op.SetDecisions(DecisionsAttr::get(
+        $_op.SetDecisions(0, DecisionsAttr::get(
           decisions.sequence(), decisions.loop_nest(), new_storage,
           decisions.expansion(), $_op.getContext()));
       }]
     >,
     InterfaceMethod<
       "Returns the sequence number of this operation as signed integer.",
       "::llvm::Optional<int64_t>", "Sequence", (ins), [{}], [{
-        if (!$_op.decisions().hasValue()) return llvm::None;
-        if ($_op.decisions()->sequence() == nullptr) return llvm::None;
-        return $_op.decisions()->sequence().getInt();
+        if ($_op.NumInstances() < 1) return llvm::None;
+        DecisionsAttr decisions = $_op.GetDecisions(0);
+        if (decisions.sequence() == nullptr) return ::llvm::None;
+        return decisions.sequence().getInt();
       }]
     >,
     InterfaceMethod<
       "Sets the sequence number of the operation to the given value.",
       "void", "SetSequence", (ins "int64_t":$seq), [{}], [{
         ::mlir::Builder builder($_op->getContext());
-        DecisionsAttr decisions = GetDecisions();
-        $_op.SetDecisions(DecisionsAttr::get(
+        DecisionsAttr decisions = GetDecisions(0);
+        $_op.SetDecisions(0, DecisionsAttr::get(
           builder.getI64IntegerAttr(seq), decisions.loop_nest(),
           decisions.storage(), decisions.expansion(),
           $_op.getContext()));
@@ -511,7 +544,7 @@ def SairComputeOp : OpInterface<"ComputeOp"> {
   let verify = [{return sair::VerifyComputeOp(op);}];
 
   let extraClassDeclaration = [{
-    static constexpr llvm::StringRef kDecisionsAttrName = "decisions";
+    static constexpr llvm::StringRef kInstancesAttrName = "instances";
   }];
 }