Add support for hoisting bounded alloc ops (#16171)

iree-org · Jan 20, 2024 · 3998b2d · 3998b2d
1 parent 04d5849
commit 3998b2d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/HoistStaticallyBoundAllocations.cpp b/compiler/src/iree/compiler/Codegen/Common/HoistStaticallyBoundAllocations.cpp
@@ -29,6 +29,7 @@ void HoistStaticallyBoundAllocationsPass::runOnOperation() {
   func::FuncOp funcOp = getOperation();
   IRRewriter rewriter(funcOp->getContext());
   hoistStaticallyBoundAllocationsInFunc<memref::AllocaOp>(rewriter, funcOp);
+  hoistStaticallyBoundAllocationsInFunc<memref::AllocOp>(rewriter, funcOp);
 }
 
 std::unique_ptr<OperationPass<func::FuncOp>>

diff --git a/compiler/src/iree/compiler/Codegen/Common/test/hoist_statically_bound_allocations.mlir b/compiler/src/iree/compiler/Codegen/Common/test/hoist_statically_bound_allocations.mlir
@@ -128,3 +128,28 @@ func.func @nested_op_alloca_yields(%arg0 : index, %arg1 : memref<?xi32>) -> memr
 //       CHECK:   scf.for
 //       CHECK:     %[[ALLOCA:.+]] = memref.alloca
 //       CHECK:     scf.yield %[[ALLOCA]]
+
+// -----
+
+#map = affine_map<(d0) -> (d0, 16)>
+func.func @nested_op_alloc_linalg_use(%arg0 : index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c42 = arith.constant 42 : i32
+  scf.for %iv = %c0 to %arg0 step %c1 {
+    %0 = affine.min #map(%iv)
+    %1 = memref.alloc(%0) : memref<?xi32>
+    linalg.fill ins(%c42 : i32) outs(%1 : memref<?xi32>)
+    memref.dealloc %1 : memref<?xi32>
+    scf.yield
+  }
+  return
+}
+// CHECK-LABEL: func @nested_op_alloc_linalg_use(
+//  CHECK-NEXT:   %[[ALLOC:.+]] = memref.alloc() : memref<16xi32>
+//       CHECK:   scf.for
+//       CHECK:     %[[SIZE:.+]] = affine.min
+//       CHECK:     %[[SUBVIEW:.+]] = memref.subview %[[ALLOC]][0] [%[[SIZE]]] [1]
+//       CHECK:     linalg.fill
+//  CHECK-SAME:         outs(%[[SUBVIEW]] :
+//       CHECK:   memref.dealloc %[[ALLOC:.+]] : memref<16xi32>