[mlir] alloc-to-alloca conversion for memref (#65335)

Introduce a simple conversion of a memref.alloc/dealloc pair into an
alloca in the same scope. Expose it as a transform op and a pattern.

Allocas typically lower to stack allocations as opposed to alloc/dealloc
that lower to significantly more expensive malloc/free calls. In
addition, this can be combined with allocation hoisting from loops to
further improve performance.

Author: ftynse

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

@@ -51,6 +51,23 @@ def MemrefToLLVMTypeConverterOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyAllocToAllocaOp : Op<Transform_Dialect,
+    "apply_patterns.memref.alloc_to_alloca",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface, ["populatePatternsWithState"]>]> {
+  let description = [{
+    Collects patterns to rewrite scoped dynamic allocation (`alloc`/`dealloc`
+    pairs) into automatic allocation (`alloca`) in the same scope, for memrefs
+    of static shape.
+
+    The `size_limit` attribute controls the maximum allocated memory (in bytes,
+    subject to data layout) for which the pattern applies.
+  }];
+
+  let arguments = (ins
+    OptionalAttr<I64Attr>:$size_limit);
+  let assemblyFormat = "(`size_limit` `(` $size_limit^ `)`)? attr-dict";
+}
+
 def ApplyExpandOpsPatternsOp : Op<Transform_Dialect,
     "apply_patterns.memref.expand_ops",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 
 namespace mlir {
 class OpBuilder;
@@ -31,6 +32,7 @@ class NarrowTypeEmulationConverter;
 namespace memref {
 class AllocOp;
 class AllocaOp;
+class DeallocOp;
 
 //===----------------------------------------------------------------------===//
 // Patterns
@@ -196,6 +198,15 @@ FailureOr<Value> replaceWithIndependentOp(RewriterBase &rewriter,
                                           memref::AllocaOp allocaOp,
                                           ValueRange independencies);
 
+/// Replaces the given `alloc` with the corresponding `alloca` and returns it if
+/// the following conditions are met:
+///   - the corresponding dealloc is available in the same block as the alloc;
+///   - the filter, if provided, succeeds on the alloc/dealloc pair.
+/// Otherwise returns nullptr and leaves the IR unchanged.
+memref::AllocaOp allocToAlloca(
+    RewriterBase &rewriter, memref::AllocOp alloc,
+    function_ref<bool(memref::AllocOp, memref::DeallocOp)> filter = nullptr);
+
 } // namespace memref
 } // namespace mlir
 

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

@@ -260,6 +260,18 @@ def PatternDescriptorOpInterface : OpInterface<"PatternDescriptorOpInterface"> {
       /*name=*/"populatePatterns",
       /*arguments=*/(ins "::mlir::RewritePatternSet &":$patterns)
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Populate rewrite patterns into the given pattern set taking into account
+        the transform state.
+      }],
+      /*returnType=*/"void",
+      /*name=*/"populatePatternsWithState",
+      /*arguments=*/(ins "::mlir::RewritePatternSet &":$patterns,
+                         "::mlir::transform::TransformState &":$state),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ $_op.populatePatterns(patterns); }]
+    >
   ];
 }
 

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 
+#include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -64,6 +65,42 @@ StringRef transform::MemrefToLLVMTypeConverterOp::getTypeConverterType() {
 // Apply...PatternsOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+class AllocToAllocaPattern : public OpRewritePattern<memref::AllocOp> {
+public:
+  explicit AllocToAllocaPattern(Operation *analysisRoot, int64_t maxSize = 0)
+      : OpRewritePattern<memref::AllocOp>(analysisRoot->getContext()),
+        dataLayoutAnalysis(analysisRoot), maxSize(maxSize) {}
+
+  LogicalResult matchAndRewrite(memref::AllocOp op,
+                                PatternRewriter &rewriter) const override {
+    return success(memref::allocToAlloca(
+        rewriter, op, [this](memref::AllocOp alloc, memref::DeallocOp dealloc) {
+          MemRefType type = alloc.getMemref().getType();
+          if (!type.hasStaticShape())
+            return false;
+
+          const DataLayout &dataLayout = dataLayoutAnalysis.getAtOrAbove(alloc);
+          int64_t elementSize = dataLayout.getTypeSize(type.getElementType());
+          return maxSize == 0 || type.getNumElements() * elementSize < maxSize;
+        }));
+  }
+
+private:
+  DataLayoutAnalysis dataLayoutAnalysis;
+  int64_t maxSize;
+};
+} // namespace
+
+void transform::ApplyAllocToAllocaOp::populatePatterns(
+    RewritePatternSet &patterns) {}
+
+void transform::ApplyAllocToAllocaOp::populatePatternsWithState(
+    RewritePatternSet &patterns, transform::TransformState &state) {
+  patterns.insert<AllocToAllocaPattern>(
+      state.getTopLevel(), static_cast<int64_t>(getSizeLimit().value_or(0)));
+}
+
 void transform::ApplyExpandOpsPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   memref::populateExpandOpsPatterns(patterns);

mlir/lib/Dialect/MemRef/Transforms/IndependenceTransforms.cpp

@@ -178,3 +178,27 @@ FailureOr<Value> memref::replaceWithIndependentOp(RewriterBase &rewriter,
                                 replacement->getDefiningOp());
   return replacement;
 }
+
+memref::AllocaOp memref::allocToAlloca(
+    RewriterBase &rewriter, memref::AllocOp alloc,
+    function_ref<bool(memref::AllocOp, memref::DeallocOp)> filter) {
+  memref::DeallocOp dealloc = nullptr;
+  for (Operation &candidate :
+       llvm::make_range(alloc->getIterator(), alloc->getBlock()->end())) {
+    dealloc = dyn_cast<memref::DeallocOp>(candidate);
+    if (dealloc && dealloc.getMemref() == alloc.getMemref() &&
+        (!filter || filter(alloc, dealloc))) {
+      break;
+    }
+  }
+
+  if (!dealloc)
+    return nullptr;
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(alloc);
+  auto alloca = rewriter.replaceOpWithNewOp<memref::AllocaOp>(
+      alloc, alloc.getMemref().getType(), alloc.getOperands());
+  rewriter.eraseOp(dealloc);
+  return alloca;
+}

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

@@ -378,8 +378,8 @@ DiagnosedSilenceableFailure transform::ApplyPatternsOp::applyToOne(
   RewritePatternSet patterns(ctx);
   if (!getRegion().empty()) {
     for (Operation &op : getRegion().front()) {
-      cast<transform::PatternDescriptorOpInterface>(&op).populatePatterns(
-          patterns);
+      cast<transform::PatternDescriptorOpInterface>(&op)
+          .populatePatternsWithState(patterns, state);
     }
   }
 

mlir/test/Dialect/MemRef/alloc-to-alloca.mlir

@@ -0,0 +1,68 @@
+// RUN: mlir-opt --test-transform-dialect-interpreter=debug-transform-root-tag=all %s | FileCheck %s --check-prefixes=CHECK,ALL
+// RUN: mlir-opt --test-transform-dialect-interpreter=debug-transform-root-tag=small %s | FileCheck %s --check-prefixes=CHECK,SMALL
+
+func.func private @callee(memref<*xf32>)
+
+// CHECK-LABEL: @large_alloc
+func.func @large_alloc() {
+  // SMALL: memref.alloc()
+  // ALL:   memref.alloca
+  %0 = memref.alloc() : memref<100x100xf32>
+  %1 = memref.cast %0 : memref<100x100xf32> to memref<*xf32>
+  call @callee(%1) : (memref<*xf32>) -> ()
+  // SMALL: memref.dealloc
+  // ALL-NOT: memref.dealloc
+  memref.dealloc %0 : memref<100x100xf32>
+  return
+}
+
+// CHECK-LABEL: @small_alloc
+func.func @small_alloc() {
+  // CHECK: memref.alloca
+  %0 = memref.alloc() : memref<2x2xf32>
+  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
+  call @callee(%1) : (memref<*xf32>) -> ()
+  // CHECK-NOT: memref.dealloc
+  memref.dealloc %0 : memref<2x2xf32>
+  return
+}
+
+// CHECK-LABEL: @no_dealloc
+func.func @no_dealloc() {
+  // CHECK: memref.alloc()
+  %0 = memref.alloc() : memref<2x2xf32>
+  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
+  call @callee(%1) : (memref<*xf32>) -> ()
+  return
+}
+
+// CHECK-LABEL: @mismatching_scope
+func.func @mismatching_scope() {
+  // CHECK: memref.alloc()
+  %0 = memref.alloc() : memref<2x2xf32>
+  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
+  call @callee(%1) : (memref<*xf32>) -> ()
+  scf.execute_region {
+    memref.dealloc %0 : memref<2x2xf32>
+    scf.yield
+  }
+  return
+}
+
+transform.sequence failures(propagate) attributes {transform.target_tag = "all"} {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns to %0 {
+    transform.apply_patterns.memref.alloc_to_alloca
+  } : !transform.any_op
+  transform.yield
+}
+
+transform.sequence failures(propagate) attributes {transform.target_tag = "small"} {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns to %0 {
+    transform.apply_patterns.memref.alloc_to_alloca size_limit(32)
+  } : !transform.any_op
+  transform.yield
+}