Materialize buffers pass.

Add a pass that introduces alloc and free operations for buffers.

PiperOrigin-RevId: 361105497

Author: ulysseB

loop_nest.cc

@@ -789,9 +789,26 @@ LoopNest LoopFusionAnalysis::GetLoopNest(
   }
   result.domain_to_loops =
       MappingAttr::get(context_, result.domain.size(), iters);
-  result.shape = DomainShapeAttr::get(context_, shape_dims);
+  result.normalized_shape = DomainShapeAttr::get(context_, shape_dims);
 
   return result;
 }
 
+DomainShapeAttr LoopNest::DomainShape() const {
+  llvm::SmallVector<DomainShapeDim> shape_dims;
+  shape_dims.reserve(domain.size());
+  MappingAttr loops_to_domain = domain_to_loops.Inverse();
+  for (const ValueAccess &access : domain) {
+    MappingAttr mapping = loops_to_domain.Resize(shape_dims.size())
+                              .Inverse()
+                              .Compose(access.mapping);
+    shape_dims.emplace_back(access.value.getType().cast<RangeType>(), mapping);
+  }
+  return DomainShapeAttr::get(normalized_shape.getContext(), shape_dims);
+}
+
+DomainShapeAttr LoopNest::Shape() const {
+  return DomainShape().AccessedShape(domain_to_loops);
+}
+
 }  // namespace sair

loop_nest.h

@@ -111,10 +111,19 @@ struct LoopFusionClass {
 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 shape;
+
+  // Shape of the nest, normalized so that dependencies between dimensions are
+  // identity mappings.
+  DomainShapeAttr normalized_shape;
+
+  // Shape of the domain the loop nest is defined from.
+  DomainShapeAttr DomainShape() const;
+
+  // Shape of the loop nest.
+  DomainShapeAttr Shape() const;
 };
 
 // Computes loop fusion classes in a sair program.

storage.cc

@@ -52,7 +52,8 @@ static llvm::SmallVector<ValueViewOp, 4> SortValueViews(SairProgramOp program) {
   return sorted;
 }
 
-Buffer::Buffer(mlir::Type element_type, int rank, ComputeOp op, int result)
+Buffer::Buffer(mlir::Type element_type, int rank, ComputeOp op, int result,
+               const LoopFusionAnalysis &fusion_analysis)
     : loc_(op.getLoc()), element_type_(element_type) {
   assert(element_type != nullptr);
   assert(rank >= 0);
@@ -67,28 +68,28 @@ Buffer::Buffer(mlir::Type element_type, int rank, ComputeOp op, int result)
   auto none_expr = MappingNoneExpr::get(element_type.getContext());
   layout_.resize(rank, none_expr);
   writes_.emplace_back(op, result);
+
+  loop_nest_mapping_ = fusion_analysis.GetLoopNest(loop_nest_).domain_to_loops;
 }
 
-void Buffer::TrimLoopNest(int new_size,
-                          const LoopFusionAnalysis &fusion_analysis) {
+void Buffer::TrimLoopNest(int new_size) {
   assert(new_size <= loop_nest_.size());
   loop_nest_.resize(new_size);
+  loop_nest_mapping_ = loop_nest_mapping_.Resize(new_size);
   if (domain_.empty()) return;
 
   mlir::MLIRContext *context = element_type_.getContext();
   // Compute dimensions used by layout.
-  llvm::SmallBitVector used_dimensions(domain_.size());
+  llvm::SmallBitVector used_dimensions = loop_nest_mapping_.DependencyMask();
   for (MappingExpr layout_expr : layout_) {
     layout_expr.SetDependenciesInMask(used_dimensions);
   }
 
   // Trim domain from unused dimensions.
   llvm::SmallVector<ValueAccess> old_domain;
   std::swap(old_domain, domain_);
-  llvm::append_range(domain_, fusion_analysis.GetLoopNest(loop_nest_).domain);
   llvm::SmallVector<MappingExpr> renaming(old_domain.size(),
                                           MappingNoneExpr::get(context));
-
   for (int dim : used_dimensions.set_bits()) {
     // Already added to the new domain.
     if (renaming[dim].isa<MappingDimExpr>()) continue;
@@ -114,6 +115,19 @@ void Buffer::AddWrite(ComputeOp op, int result) {
   writes_.emplace_back(op, result);
 }
 
+void Buffer::AddRead(ComputeOp op, int operand) {
+  reads_.emplace_back(op, operand);
+}
+
+MappingAttr Buffer::PrefixedLayout() const {
+  mlir::MLIRContext *context = loop_nest_mapping_.getContext();
+  llvm::SmallVector<MappingExpr> exprs;
+  exprs.reserve(loop_nest_.size() + layout_.size());
+  llvm::append_range(exprs, loop_nest_mapping_);
+  llvm::append_range(exprs, layout_);
+  return MappingAttr::get(context, domain_.size(), exprs);
+}
+
 StorageAnalysis::StorageAnalysis(mlir::Operation *operation)
     : StorageAnalysis(operation->getContext()) {
   mlir::LogicalResult result = Init(cast<SairProgramOp>(operation));
@@ -234,7 +248,8 @@ static mlir::LogicalResult DeclareBuffer(
       op->getResult(result).getType().cast<ValueType>().ElementType();
 
   int rank = attr.layout().mapping().size();
-  auto it = buffer_map.try_emplace(attr.name(), element_type, rank, op, result);
+  auto it = buffer_map.try_emplace(attr.name(), element_type, rank, op, result,
+                                   fusion_analysis);
   Buffer &buffer = it.first->second;
 
   if (!it.second) {
@@ -281,7 +296,7 @@ static mlir::LogicalResult DeclareBuffer(
     }
   }
 
-  buffer.TrimLoopNest(num_deps, fusion_analysis);
+  buffer.TrimLoopNest(num_deps);
   return mlir::success();
 }
 
@@ -408,20 +423,13 @@ static mlir::LogicalResult CheckMallocInsertionPoint(
 // `loop_nest`. Increases `min_num_loops` to the minimal number of loops needed
 // in `loop_nest` for the layout to be valid.
 static mlir::LogicalResult CheckLayoutMapping(
-    const LoopNest &loop_nest, mlir::StringAttr buffer_name,
-    const Buffer &buffer, const IterationSpaceAnalysis &iteration_spaces,
-    int &min_num_loops) {
-  mlir::MLIRContext *context = loop_nest.domain_to_loops.getContext();
+    mlir::StringAttr buffer_name, const Buffer &buffer,
+    const IterationSpaceAnalysis &iteration_spaces, int &min_num_loops) {
+  mlir::MLIRContext *context = buffer_name.getContext();
   int domain_size = buffer.domain().size();
   int loop_nest_size = buffer.loop_nest().size();
 
-  // Get a mapping that maps both loop-nest and layout indices. This corresponds
-  // to the different instances of the buffer.
-  llvm::SmallVector<MappingExpr> exprs;
-  exprs.reserve(loop_nest_size + buffer.layout().size());
-  llvm::append_range(exprs, loop_nest.domain_to_loops);
-  llvm::append_range(exprs, buffer.layout());
-  auto mapping = MappingAttr::get(context, domain_size, exprs);
+  MappingAttr mapping = buffer.PrefixedLayout();
 
   // Update `min_num_loops` based on domain dimensions layout depends on.
   llvm::SmallBitVector used_dimensions(domain_size);
@@ -499,13 +507,13 @@ mlir::LogicalResult StorageAnalysis::Init(SairProgramOp program) {
     // Check that layout mapping is correct and compute the minimal loop nest
     // each buffer needs to be nested in.
     int min_num_loops = 0;
-    if (mlir::failed(CheckLayoutMapping(loop_nest, name, buffer,
-                                        iteration_spaces, min_num_loops))) {
+    if (mlir::failed(CheckLayoutMapping(name, buffer, iteration_spaces,
+                                        min_num_loops))) {
       return mlir::failure();
     }
 
     // Minimize layout loop-nest.
-    buffer.TrimLoopNest(min_num_loops, fusion_analysis);
+    buffer.TrimLoopNest(min_num_loops);
   }
 
   // Compute value storages.
@@ -529,6 +537,16 @@ mlir::LogicalResult StorageAnalysis::Init(SairProgramOp program) {
     }
   }
 
+  // Register buffer reads.
+  program.walk([&](ComputeOp op) {
+    for (ValueOperand operand : cast<SairOp>(*op).ValueOperands()) {
+      const ValueStorage &storage = value_storages_[operand.value()];
+      if (storage.buffer_name == nullptr) continue;
+      buffers_.find(storage.buffer_name)
+          ->second.AddRead(op, operand.position());
+    }
+  });
+
   return mlir::success();
 }
 

storage.h

@@ -36,7 +36,8 @@ class Buffer {
   // Create a new buffer written to by the given operation. The operation must
   // have a loop_nest attribute set. `result` is the position of `op` result
   // stored in `buffer`.
-  Buffer(mlir::Type element_type, int rank, ComputeOp op, int result);
+  Buffer(mlir::Type element_type, int rank, ComputeOp op, int result,
+         const LoopFusionAnalysis &fusion_analysis);
 
   // Number of dimensions in the buffer layout.
   int rank() const { return layout_.size(); }
@@ -58,25 +59,40 @@ class Buffer {
   // result stored in the buffer. Never empty.
   llvm::ArrayRef<std::pair<ComputeOp, int>> writes() const { return writes_; }
 
+  // List of operations that read from the buffer, with the position of the Sair
+  // value operand.
+  llvm::ArrayRef<std::pair<ComputeOp, int>> reads() const { return reads_; }
+
   // Get the location of the first operation defining the buffer.
   mlir::Location getLoc() const { return loc_; }
 
+  // Mapping of domain to layout prefixed by loop nest iterators. The prefix
+  // corresponds to the different instances of the buffer.
+  MappingAttr PrefixedLayout() const;
+
   // Registers an operation writting to the buffer.
   void AddWrite(ComputeOp op, int result);
 
+  // Registers an operation reading the buffer.
+  void AddRead(ComputeOp op, int operand);
+
   // Trims the loop-nest to the given size.
-  void TrimLoopNest(int new_size, const LoopFusionAnalysis &fusion_analysis);
+  void TrimLoopNest(int new_size);
 
   // Unifies a dimension of the layout with another expression.
   void UnifyLayoutDim(int layout_dim, MappingExpr expr);
 
  private:
   mlir::Location loc_;
   mlir::Type element_type_;
+
   llvm::SmallVector<mlir::StringAttr> loop_nest_;
+  MappingAttr loop_nest_mapping_;
+
   llvm::SmallVector<ValueAccess> domain_;
   llvm::SmallVector<MappingExpr> layout_;
   llvm::SmallVector<std::pair<ComputeOp, int>> writes_;
+  llvm::SmallVector<std::pair<ComputeOp, int>> reads_;
 };
 
 // Computes buffers metadata and storage information for each value.
@@ -95,6 +111,11 @@ class StorageAnalysis {
     return buffers_.find(buffer)->second;
   }
 
+  // List of buffers indexed by name.
+  const llvm::DenseMap<mlir::Attribute, Buffer> &buffers() const {
+    return buffers_;
+  }
+
   // Retrieves the storage of a value.
   const ValueStorage &GetStorage(mlir::Value value) const {
     return value_storages_.find(value)->second;

test/materialize_buffers.mlir

@@ -0,0 +1,115 @@
+// RUN: sair-opt %s -sair-materialize-buffers -mlir-print-local-scope | FileCheck %s
+
+// CHECK-LABEL: @static_shape
+func @static_shape(%arg0: f32) {
+  sair.program {
+    %0 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    %1 = sair.static_range 16 step 2: !sair.range
+    // CHECK: %[[V0:.*]] = sair.alloc {loop_nest = [],
+    // CHECK-SAME: storage = [{layout = #sair.named_mapping<[] -> ()>, space = "register"}]
+    // CHECK-SAME: : !sair.value<(), memref<8xf32>>
+    // CHECK: sair.copy[d0:%{{.*}}]
+    %2 = sair.copy[d0:%1] %0 {
+      loop_nest = [{name = "A", iter = #sair.mapping_expr<d0>}],
+      storage = [{
+        name = "B", space = "memory",
+        layout = #sair.named_mapping<[d0:"A"] -> (d0)>
+      }]
+    } : !sair.value<d0:range, f32>
+    // CHECK: sair.copy[d0:%{{.*}}]
+    %3 = sair.copy[d0:%1] %2(d0) {
+      loop_nest = [{name = "B", iter = #sair.mapping_expr<d0>}],
+      storage = [{space = "register", layout = #sair.named_mapping<[] -> ()>}]
+    } : !sair.value<d0:range, f32>
+    // CHECK: sair.free %[[V0]] {loop_nest = []} : !sair.value<(), memref<8xf32>>
+    sair.exit
+  }
+  return
+}
+
+// CHECK-LABEL: @dynamic_shape
+func @dynamic_shape(%arg0: f32, %arg1: index, %arg2: index) {
+  sair.program {
+    %0 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    // CHECK: %[[V1:.*]] = sair.from_scalar %{{.*}} : !sair.value<(), index>
+    %1 = sair.from_scalar %arg1 : !sair.value<(), index>
+    // CHECK: %[[V2:.*]] = sair.from_scalar %{{.*}} : !sair.value<(), index>
+    %2 = sair.from_scalar %arg2 : !sair.value<(), index>
+    %3 = sair.dyn_range %1, %2 step 4 : !sair.range
+    // CHECK: %[[V3:.*]] = sair.map %[[V1]], %[[V2]] attributes {
+    // CHECK:   loop_nest = []
+    // CHECK:   storage = [{layout = #sair.named_mapping<[] -> ()>, space = "register"}]
+    // CHECK: } {
+    // CHECK:   ^{{.*}}(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index):
+    // CHECK:     %[[V4:.*]] = affine.apply
+    // CHECK:       affine_map<(d0, d1) -> ((d1 - d0) ceildiv 4)>(%[[ARG0]], %[[ARG1]])
+    // CHECK:     sair.return %[[V4]]
+    // CHECK: } : #sair.shape<()>, (index, index) -> index
+
+    // CHECK: %[[V5:.*]] = sair.alloc %[[V3]]
+    // CHECK:   : !sair.value<(), memref<?xf32>>
+    %4 = sair.copy[d0:%3] %0 {
+      loop_nest = [{name = "A", iter = #sair.mapping_expr<d0>}],
+      storage = [{
+        name = "B", space = "memory",
+        layout = #sair.named_mapping<[d0:"A"] -> (d0)>
+      }]
+    } : !sair.value<d0:range, f32>
+    sair.exit
+  }
+  return
+}
+
+// CHECK-LABEL: @loop_nest
+func @loop_nest(%arg0: f32) {
+  sair.program {
+    %0 = sair.from_scalar %arg0 : !sair.value<(), f32>
+    %1 = sair.static_range 16 : !sair.range
+    // CHECK: %[[D0:.*]] = sair.placeholder : !sair.range
+
+    // CHECK: %[[V0:.*]] = sair.map[d0:%[[D0]]] attributes {
+    // CHECK:   loop_nest = [{iter = #sair.mapping_expr<d0>, name = "A"}]
+    // CHECK: } {
+    // CHECK:   ^{{.*}}(%[[ARG0:.*]]: index):
+    // CHECK:     %[[V1:.*]] = affine.apply affine_map<(d0) -> (d0)>(%[[ARG0]])
+    // CHECK:     %[[C4:.*]] = constant 4
+    // CHECK:     %[[V2:.*]] = addi %[[V1]], %[[C4]]
+    // CHECK:     %[[C16:.*]] = constant 16
+    // CHECK:     %[[V3:.*]] = cmpi ult, %[[C16]], %[[V2]]
+    // CHECK:     %[[V4:.*]] = select %[[V3]], %[[C16]], %[[V2]]
+    // CHECK:     %[[V5:.*]] = affine.apply affine_map<(d0, d1) -> (d1 - d0)>
+    // CHECK:     sair.return %[[V5]] : index
+    // CHECK: } : #sair.shape<d0:range>, () -> index
+
+    // CHECK: %[[V6:.*]] = sair.alloc[d0:%[[D0]]] %[[V0]](d0) {
+    // CHECK:   loop_nest = [{iter = #sair.mapping_expr<d0>, name = "A"}]
+    // CHECK: }  : !sair.value<d0:range, memref<?xf32>>
+    // CHECK: sair.copy
+    %2 = sair.copy[d0:%1] %0 {
+      loop_nest = [
+        {name = "A", iter = #sair.mapping_expr<stripe(d0, 4)>},
+        {name = "B", iter = #sair.mapping_expr<stripe(d0, 1 size 4)>}
+      ],
+      storage = [{
+        name = "buf", space = "memory",
+        layout = #sair.named_mapping<[d0:"B"] -> (d0)>
+      }]
+    } : !sair.value<d0:range, f32>
+    // CHECK: sair.copy
+    %3 = sair.copy[d0:%1] %2(d0) {
+      loop_nest = [
+        {name = "A", iter = #sair.mapping_expr<stripe(d0, 4)>},
+        {name = "C", iter = #sair.mapping_expr<stripe(d0, 1 size 4)>}
+      ],
+      storage = [{
+        name = "buf", space = "memory",
+        layout = #sair.named_mapping<[d0:"C"] -> (d0)>
+      }]
+    } : !sair.value<d0:range, f32>
+    // CHECK: sair.free[d0:%[[D0]]] %[[V6]](d0) {
+    // CHECK:   loop_nest = [{iter = #sair.mapping_expr<d0>, name = "A"}]
+    // CHECK: } : !sair.value<d0:range, memref<?xf32>>
+    sair.exit
+  }
+  return
+}

transforms/CMakeLists.txt

@@ -65,6 +65,7 @@ add_mlir_library(sair_lowering
   introduce_loops.cc
   introduce_memrefs.cc
   lower_to_map.cc
+  materialize_buffers.cc
   normalize_loops.cc
 
   DEPENDS

transforms/lowering.h

@@ -55,6 +55,10 @@ std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> CreateInlineTrivialOpsPass();
 std::unique_ptr<mlir::OperationPass<mlir::FuncOp>>
 CreateMaterializeMemRefsPass();
 
+// Replaces Sair values by buffers as specified by the `storage` attribute.
+std::unique_ptr<mlir::OperationPass<mlir::FuncOp>>
+CreateMaterializeBuffersPass();
+
 // Replaces iteration dimensions by loops in sair.map and sair.map_reduce
 // operations.
 std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> CreateIntroduceLoopsPass();