Replace memory_space attribute by storage attribute.

The `storage` attribute applies to `ComputeOp` operations only. It specifies an optional `BufferAttr` attribute for each output. The buffer attribute specifies:
* A memory space, currently equal to "memory" or "register"
* A name, for "memory" buffers only. All buffers sharing the same name are stored in the same memref.
* a mapping from a subset of loops to dimensions of the memref backing the buffer.

* Remove `memory_space` attribute from all operations.
* Add a `storage` attribute to ComputeOp operations.
* Introduce a `ValueViewOp` interface to infer the storage of non-compute operations.
* Extend `LoopFusionAnalysis` to return the domain and the mapping from domain to loop iteration of loop nest.
* Define a `StorageAnalysis` that checks and unify buffer attributes with the same names.

As a side effect of change `memory_space` into `storage`:
* Update passes using `memory_space` to use `storage` instead.
* Change DefaultMemorySpace pass into a DefaultStorage pass. Only support the subset of operations that was previously supported by DefaultMemorySpace.
* Update sair-from-linalg pass to insert copy operations after sair.from_memerf operations. This clears the storage of values originating from sair.from_memref so that is does not conflicts with other storage constraints.
* Remove unused ValueProducer interface.
* Update ResolveUnificationConstraints in `util.cc` so that it can also handle unification of buffers layouts.

PiperOrigin-RevId: 361096395

Author: ulysseB

CMakeLists.txt

@@ -83,6 +83,7 @@ add_mlir_library(sair_dialect
   sair_ops.cc
   sair_types.cc
   util.cc
+  storage.cc
 
   DEPENDS
   sair_ops_inc_gen

canonicalization_patterns.cc

@@ -77,7 +77,7 @@ bool SimplifyProjOp(ValueOperand &use, ProjOp op,
   rewriter.setInsertionPoint(op);
   ProjOp new_op = rewriter.create<ProjOp>(
       op.getLoc(), op.getType(), op.parallel_domain(), projection_domain,
-      mapping_array, prev_op.value(), shape, op.memory_spaceAttr());
+      mapping_array, prev_op.value(), shape);
   use.set_value(new_op.result());
 
   return true;
@@ -141,7 +141,7 @@ class DeduplicateMapInputsOutputs : public OpRewritePattern<SairMapOp> {
     llvm::SmallVector<mlir::Value, 4> old_results_to_keep;
     llvm::SmallVector<mlir::Value, 4> new_scalar_results;
     llvm::SmallVector<mlir::Type, 4> new_result_types;
-    llvm::SmallVector<mlir::Attribute, 4> new_memory_spaces;
+    llvm::SmallVector<mlir::Attribute, 4> new_storages;
 
     for (ValueOperand operand : op.ValueOperands()) {
       mlir::Value argument =
@@ -176,7 +176,7 @@ class DeduplicateMapInputsOutputs : public OpRewritePattern<SairMapOp> {
       // Deduplicate results.
       for (int j = 0; j < i; ++j) {
         if (scalar_value != return_op.getOperand(j)) continue;
-        if (op.GetMemorySpace(i) != op.GetMemorySpace(j)) continue;
+        if (op.Storage(i) != op.Storage(j)) continue;
         result.replaceAllUsesWith(op.getResult(j));
         break;
       }
@@ -187,13 +187,9 @@ class DeduplicateMapInputsOutputs : public OpRewritePattern<SairMapOp> {
       old_results_to_keep.push_back(result);
       new_scalar_results.push_back(scalar_value);
       new_result_types.push_back(result.getType());
-      mlir::Attribute memory_space =
-          op.GetMemorySpace(i)
-              .map([&](int value) -> mlir::Attribute {
-                return rewriter.getI32IntegerAttr(value);
-              })
-              .getValueOr(rewriter.getUnitAttr());
-      new_memory_spaces.push_back(memory_space);
+      mlir::Attribute new_storage = op.Storage(i);
+      new_storages.push_back(new_storage == nullptr ? rewriter.getUnitAttr()
+                                                    : new_storage);
     }
 
     // Create the new operation if necessary.
@@ -210,7 +206,7 @@ class DeduplicateMapInputsOutputs : public OpRewritePattern<SairMapOp> {
     SairMapOp new_op = rewriter.create<SairMapOp>(
         op.getLoc(), new_result_types, op.domain(),
         rewriter.getArrayAttr(new_mappings), new_operands, op.shape(),
-        op.loop_nestAttr(), rewriter.getArrayAttr(new_memory_spaces));
+        op.loop_nestAttr(), rewriter.getArrayAttr(new_storages));
     new_op.body().takeBody(op.body());
 
     for (auto p : llvm::zip(old_results_to_keep, new_op.results())) {
@@ -310,7 +306,7 @@ class RemoveUnreferencedDims : public OpRewritePattern<OpTy> {
         op.getType().template cast<ValueType>().AccessedType(partial_mapping),
         parallel_dimensions, projection_dimensions,
         rewriter.getArrayAttr(mapping.Inverse().Compose(op.Value().Mapping())),
-        op.value(), op.shape().AccessedShape(mapping), op.memory_spaceAttr());
+        op.value(), op.shape().AccessedShape(mapping));
     new_op->setDialectAttrs(op->getDialectAttrs());
 
     // Replace the original op.
@@ -356,7 +352,7 @@ class RemoveUnreferencedDims<SairFbyOp> : public OpRewritePattern<SairFbyOp> {
         parallel_dimensions, sequential_dimensions,
         rewriter.getArrayAttr({inverted_mapping.Compose(op.Init().Mapping()),
                                inverted_mapping.Compose(op.Value().Mapping())}),
-        op.init(), op.value(), op.memory_spaceAttr());
+        op.init(), op.value());
     new_op->setDialectAttrs(op->getDialectAttrs());
 
     // Replace the original op.

loop_nest.cc

@@ -496,30 +496,8 @@ static mlir::LogicalResult VerifySubDomains(
 }
 
 mlir::LogicalResult VerifyLoopNests(SairProgramOp program) {
-  // Verify operands of Sair operands are defined in the same program. This
-  // check is performed here rather that in SairOp as it is needed for other
-  // verifications.
-  mlir::WalkResult result = program.walk([&](SairOp op) -> mlir::WalkResult {
-    for (mlir::Value dimension : op.domain()) {
-      mlir::Operation *defining_op = dimension.getDefiningOp();
-      if (defining_op == nullptr || defining_op->getParentOp() != program) {
-        return op.emitError()
-               << "sair dimensions must be defined in the region they are used";
-      }
-    }
-    for (ValueOperand operand : op.ValueOperands()) {
-      mlir::Operation *defining_op = operand.value().getDefiningOp();
-      if (defining_op == nullptr || defining_op->getParentOp() != program) {
-        return op.emitError()
-               << "sair values must be defined in the region they are used";
-      }
-    }
-    return mlir::success();
-  });
-  if (result.wasInterrupted()) return mlir::failure();
-
   // Verify loop nests are correct with regard to their operation.
-  result = program.walk([](ComputeOp op) -> mlir::WalkResult {
+  mlir::WalkResult result = program.walk([](ComputeOp op) -> mlir::WalkResult {
     if (!op.loop_nest().hasValue()) return mlir::WalkResult::advance();
     return VerifyLoopNestWellFormed(
         cast<SairOp>(op.getOperation()), op.LoopNestLoops());
@@ -569,7 +547,8 @@ mlir::LogicalResult VerifyLoopNests(SairProgramOp program) {
   return mlir::success();
 }
 
-LoopFusionAnalysis::LoopFusionAnalysis(mlir::Operation *operation) {
+LoopFusionAnalysis::LoopFusionAnalysis(mlir::Operation *operation)
+    : context_(operation->getContext()) {
   SairProgramOp program_op = dyn_cast<SairProgramOp>(operation);
   if (program_op == nullptr) return;
   mlir::LogicalResult status = Init(program_op);
@@ -579,19 +558,17 @@ LoopFusionAnalysis::LoopFusionAnalysis(mlir::Operation *operation) {
 
 std::optional<LoopFusionAnalysis> LoopFusionAnalysis::Create(
     SairProgramOp program_op) {
-  LoopFusionAnalysis analysis;
+  LoopFusionAnalysis analysis(program_op->getContext());
   if (mlir::failed(analysis.Init(program_op))) return std::nullopt;
   return analysis;
 }
 
 mlir::LogicalResult LoopFusionAnalysis::Init(SairProgramOp program_op) {
-  mlir::MLIRContext *context = program_op.getContext();
-
   llvm::SmallVector<ComputeOp> work_list;
   program_op.walk([&](ComputeOp op) {
     auto sair_op = cast<SairOp>(op.getOperation());
     int domain_size = sair_op.domain().size();
-    auto none_expr = MappingNoneExpr::get(context);
+    auto none_expr = MappingNoneExpr::get(context_);
     op_domain_mappings_[op.getOperation()].resize(domain_size, none_expr);
     if (!op.loop_nest().hasValue()) return;
     work_list.push_back(op);
@@ -638,9 +615,9 @@ mlir::LogicalResult LoopFusionAnalysis::Init(SairProgramOp program_op) {
 
     int domain_size = fusion_class.domain.size();
     MappingAttr inverse_loop_nest =
-        MappingAttr::get(context, domain_size, loop_nest).Inverse();
+        MappingAttr::get(context_, domain_size, loop_nest).Inverse();
 
-    auto hr_domain = DomainShapeAttr::HyperRectangular(context, domain_size);
+    auto hr_domain = DomainShapeAttr::HyperRectangular(context_, domain_size);
     DomainShapeDim loop_shape = fusion_class.iter_expr.AccessedShape(
         hr_domain.Dimensions(), inverse_loop_nest);
     if (!loop_shape.dependency_mapping().IsFullySpecified()) {
@@ -720,9 +697,23 @@ mlir::LogicalResult LoopFusionAnalysis::RegisterLoop(
   llvm::SmallBitVector constrained_dims =
       loop.iter().DependencyMask(domain_size);
   for (int dimension : constrained_dims.set_bits()) {
+    // Compute the mapping to access the dimension in from loop indices.
+    MappingAttr old_access_mapping = sair_op.shape()
+                                         .Dimension(dimension)
+                                         .dependency_mapping()
+                                         .ResizeUseDomain(domain_size);
+    MappingAttr new_access_mapping =
+        loops_to_op_domain_mapping.Compose(old_access_mapping);
+    if (!new_access_mapping.IsFullySpecified()) {
+      return op->emitError()
+             << "dimension d" << dimension << " in " << loop_name.str()
+             << " is used before its dependencies";
+    }
+
+    ValueAccess access = {sair_op.domain()[dimension], new_access_mapping};
     if (mlir::failed(ResolveUnificationConstraint(
-            op, dimension, loop_name.str(), loops_to_op_domain_mapping,
-            constraints[dimension], fusion_class.domain))) {
+            op.getLoc(), loop_name.str(), access, constraints[dimension],
+            fusion_class.domain))) {
       return mlir::failure();
     }
   }
@@ -734,4 +725,46 @@ mlir::LogicalResult LoopFusionAnalysis::RegisterLoop(
   return mlir::success();
 }
 
+LoopNest LoopFusionAnalysis::GetLoopNest(
+    llvm::ArrayRef<mlir::Attribute> loops) const {
+  llvm::SmallVector<mlir::StringAttr> loop_names;
+  loop_names.reserve(loops.size());
+  for (mlir::Attribute attr : loops) {
+    auto loop = attr.cast<LoopAttr>();
+    loop_names.push_back(loop.name());
+  }
+  return GetLoopNest(loop_names);
+}
+
+LoopNest LoopFusionAnalysis::GetLoopNest(
+    llvm::ArrayRef<mlir::StringAttr> loop_names) const {
+  LoopNest result;
+  if (!loop_names.empty()) {
+    result.domain = GetClass(loop_names.back()).domain;
+  }
+
+  llvm::SmallVector<MappingExpr> iters;
+  llvm::SmallVector<DomainShapeDim> shape_dims;
+  iters.reserve(loop_names.size());
+  shape_dims.reserve(loop_names.size());
+  for (mlir::StringAttr name : loop_names) {
+    const LoopFusionClass &fusion_class = GetClass(name);
+    iters.push_back(fusion_class.iter_expr);
+
+    // Resulting loop nest dependencies are pointwise dependencies to a prefix
+    // of the loop nest.
+    int num_dependencies = fusion_class.dependencies.size();
+    auto dim_type = RangeType::get(DomainShapeAttr::get(
+        context_, llvm::makeArrayRef(shape_dims).take_front(num_dependencies)));
+    auto dim_mapping =
+        MappingAttr::GetIdentity(context_, num_dependencies, shape_dims.size());
+    shape_dims.emplace_back(dim_type, dim_mapping);
+  }
+  result.domain_to_loops =
+      MappingAttr::get(context_, result.domain.size(), iters);
+  result.loops_shape = DomainShapeAttr::get(context_, shape_dims);
+
+  return result;
+}
+
 }  // namespace sair

loop_nest.h

@@ -23,8 +23,8 @@
 
 namespace sair {
 
-// Verifies loop nest attributes of operations nested in the
-// sair.program operation.
+// Verifies loop nest attributes of operations nested in the sair.program
+// operation. Assumes that Sair operands are defined in the same program.
 mlir::LogicalResult VerifyLoopNests(SairProgramOp program);
 
 // Analysis of how data is distributed on loop nests iterations. It indicates,
@@ -75,14 +75,14 @@ class IterationSpaceAnalysis {
 // A class of fused loops.
 struct LoopFusionClass {
   // Loops this class depends on.
-  llvm::SmallVector<mlir::StringAttr, 2> dependencies;
+  llvm::SmallVector<mlir::StringAttr> dependencies;
 
   // Domain in which the loop size is defined. This is a list of dimensions,
   // with an access pattern from dependencies indicies to the domain of each
   // dimension.
   //
   // Domains of outer fusion classes must be a prefix of this one.
-  llvm::SmallVector<ValueAccess, 4> domain;
+  llvm::SmallVector<ValueAccess> domain;
 
   // Mapping from domain indices to the loop indices.
   MappingExpr iter_expr;
@@ -91,6 +91,17 @@ struct LoopFusionClass {
   ComputeOp occurence;
 };
 
+// A loop nest of fused loops.
+// TODO: use in normalize_loops.cc
+struct LoopNest {
+  // Domain used to define loop ranges.
+  llvm::ArrayRef<ValueAccess> domain;
+  // Mapping from `domain` to loops.
+  MappingAttr domain_to_loops;
+  // Shape of the resulting loop nest.
+  DomainShapeAttr loops_shape;
+};
+
 // Computes loop fusion classes in a sair program.
 class LoopFusionAnalysis {
  public:
@@ -106,8 +117,12 @@ class LoopFusionAnalysis {
     return fusion_classes_.find(name)->second;
   }
 
+  // Retrives the unified loop nest corresponding to loops.
+  LoopNest GetLoopNest(llvm::ArrayRef<mlir::Attribute> loops) const;
+  LoopNest GetLoopNest(llvm::ArrayRef<mlir::StringAttr> loop_names) const;
+
  private:
-  LoopFusionAnalysis() {}
+  LoopFusionAnalysis(mlir::MLIRContext *context) : context_(context) {}
 
   // Populates the analysis with the operations appearing in `program_op`.
   mlir::LogicalResult Init(SairProgramOp program_op);
@@ -117,6 +132,7 @@ class LoopFusionAnalysis {
   mlir::LogicalResult RegisterLoop(ComputeOp op, LoopAttr loop,
                                    llvm::ArrayRef<mlir::Attribute> outer_loops);
 
+  mlir::MLIRContext *context_;
   llvm::DenseMap<mlir::Attribute, LoopFusionClass> fusion_classes_;
   llvm::DenseMap<mlir::Operation *, llvm::SmallVector<MappingExpr, 4>>
       op_domain_mappings_;

sair_attributes.cc

@@ -914,6 +914,14 @@ MappingAttr MappingAttr::Canonicalize() const {
   return MappingAttr::get(getContext(), UseDomainSize(), exprs);
 }
 
+int MappingAttr::MinDomainSize() const {
+  int min = 0;
+  for (MappingExpr expr : Dimensions()) {
+    min = std::max(min, expr.MinDomainSize());
+  }
+  return min;
+}
+
 //===----------------------------------------------------------------------===//
 // NamedMappingAttr
 //===----------------------------------------------------------------------===//
@@ -961,6 +969,12 @@ NamedMappingAttr NamedMappingAttr::get(llvm::ArrayRef<mlir::StringAttr> names,
   return Base::get(mapping.getContext(), names, mapping);
 }
 
+NamedMappingAttr NamedMappingAttr::GetIdentity(
+    mlir::MLIRContext *context, llvm::ArrayRef<mlir::StringAttr> names) {
+  auto mapping = MappingAttr::GetIdentity(context, names.size());
+  return NamedMappingAttr::get(names, mapping);
+}
+
 llvm::ArrayRef<mlir::StringAttr> NamedMappingAttr::names() const {
   return getImpl()->names();
 }

sair_attributes.h

@@ -123,6 +123,9 @@ class MappingAttr
   // depends on.
   llvm::SmallBitVector DependencyMask() const;
 
+  // Minimal use domain size the mapping could have while remaining valid.
+  int MinDomainSize() const;
+
   // Indicates if the mapping accesses a single element of the def
   // domain per element of the use domain, when considering only the first
   // `num_dimensions` of the use domain.
@@ -169,6 +172,10 @@ class NamedMappingAttr
   static NamedMappingAttr get(llvm::ArrayRef<mlir::StringAttr> names,
                               MappingAttr mapping);
 
+  // Constructs an instance of NamedMappingAttr with an identity mapping.
+  static NamedMappingAttr GetIdentity(mlir::MLIRContext *context,
+                                      llvm::ArrayRef<mlir::StringAttr> names);
+
   llvm::ArrayRef<mlir::StringAttr> names() const;
 
   MappingAttr mapping() const;