[mlir][Linalg] Add an option to use Alloca instead of malloc/free pairs.

Summary: A relevant test is also added.

Subscribers: mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, stephenneuendorffer, Joonsoo, grosul1, Kayjukh, jurahul, msifontes

Tags: #mlir

Differential Revision: https://reviews.llvm.org/D82959

diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.h b/mlir/include/mlir/Dialect/Linalg/Passes.h
index 8a274ed..a5c09b3 100644 (file)
--- a/mlir/include/mlir/Dialect/Linalg/Passes.h
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.h
@@ -36,7 +36,7 @@ std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
-createLinalgPromotionPass(bool dynamicBuffers);
+createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca);
 std::unique_ptr<OperationPass<FuncOp>> createLinalgPromotionPass();
 
 /// Create a pass to convert Linalg operations to scf.for loops and

diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 1fc7fa5..11f12ad 100644 (file)
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -61,7 +61,9 @@ def LinalgPromotion : FunctionPass<"linalg-promote-subviews"> {
   let constructor = "mlir::createLinalgPromotionPass()";
   let options = [
     Option<"dynamicBuffers", "test-promote-dynamic", "bool",
-           /*default=*/"false", "Test generation of dynamic promoted buffers">
+           /*default=*/"false", "Test generation of dynamic promoted buffers">,
+    Option<"useAlloca", "test-use-alloca", "bool",
+           /*default=*/"false", "Test generation of alloca'ed buffers.">
   ];
 }
 

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index b7be5d8..65f1bbb 100644 (file)
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -122,6 +122,12 @@ struct LinalgPromotionOptions {
     alignment = align;
     return *this;
   }
+  /// Use alloca with the default allocation scheme.
+  bool useAlloca = false;
+  LinalgPromotionOptions &setUseAlloca(bool use) {
+    useAlloca = use;
+    return *this;
+  }
   /// Callback function to do the allocation of the promoted buffer. If None,
   /// then the default allocation scheme of allocating a memref<?xi8> buffer
   /// followed by a view operation is used.
@@ -134,7 +140,6 @@ struct LinalgPromotionOptions {
     deallocationFn = deallocFn;
     return *this;
   }
-
   /// Callback function to do the copy of data to and from the promoted
   /// subview. If None then a linalg.copy is used.
   Optional<CopyCallbackFn> copyInFn = None;

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index de8514f..aebc33a 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -67,7 +67,8 @@ static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc,
 
 /// Alloc a new buffer of `size`. If `dynamicBuffers` is true allocate exactly
 /// the size needed, otherwise try to allocate a static bounding box.
-static Value allocBuffer(Type elementType, Value size, bool dynamicBuffers,
+static Value allocBuffer(const LinalgPromotionOptions &options,
+                         Type elementType, Value size, bool dynamicBuffers,
                          OperationFolder *folder,
                          Optional<unsigned> alignment = None) {
   auto *ctx = size.getContext();
@@ -78,23 +79,34 @@ static Value allocBuffer(Type elementType, Value size, bool dynamicBuffers,
         IntegerAttr::get(IntegerType::get(64, ctx), alignment.getValue());
   if (!dynamicBuffers)
     if (auto cst = size.getDefiningOp<ConstantIndexOp>())
-      return std_alloc(
-          MemRefType::get(width * cst.getValue(), IntegerType::get(8, ctx)),
-          ValueRange{}, alignment_attr);
+      return options.useAlloca
+                 ? std_alloca(MemRefType::get(width * cst.getValue(),
+                                              IntegerType::get(8, ctx)),
+                              ValueRange{}, alignment_attr)
+                       .value
+                 : std_alloc(MemRefType::get(width * cst.getValue(),
+                                             IntegerType::get(8, ctx)),
+                             ValueRange{}, alignment_attr)
+                       .value;
   Value mul =
       folded_std_muli(folder, folded_std_constant_index(folder, width), size);
-  return std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul,
-                   alignment_attr);
+  return options.useAlloca
+             ? std_alloca(MemRefType::get(-1, IntegerType::get(8, ctx)), mul,
+                          alignment_attr)
+                   .value
+             : std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul,
+                         alignment_attr)
+                   .value;
 }
 
 /// Default allocation callback function. This allocates a promoted buffer when
 /// no call back to do so is provided. The default is to allocate a
 /// memref<..xi8> and return a view to get a memref type of shape
 /// boundingSubViewSize.
-static Optional<Value>
-allocBufferCallBack(OpBuilder &builder, SubViewOp subView,
-                    ArrayRef<Value> boundingSubViewSize, bool dynamicBuffers,
-                    Optional<unsigned> alignment, OperationFolder *folder) {
+static Optional<Value> defaultAllocBufferCallBack(
+    const LinalgPromotionOptions &options, OpBuilder &builder,
+    SubViewOp subView, ArrayRef<Value> boundingSubViewSize, bool dynamicBuffers,
+    Optional<unsigned> alignment, OperationFolder *folder) {
   ShapedType viewType = subView.getType();
   int64_t rank = viewType.getRank();
   (void)rank;
@@ -105,7 +117,7 @@ allocBufferCallBack(OpBuilder &builder, SubViewOp subView,
   Value allocSize = one;
   for (auto size : llvm::enumerate(boundingSubViewSize))
     allocSize = folded_std_muli(folder, allocSize, size.value());
-  Value buffer = allocBuffer(viewType.getElementType(), allocSize,
+  Value buffer = allocBuffer(options, viewType.getElementType(), allocSize,
                              dynamicBuffers, folder, alignment);
   SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(),
                                    ShapedType::kDynamicSize);
@@ -118,10 +130,13 @@ allocBufferCallBack(OpBuilder &builder, SubViewOp subView,
 /// Default implementation of deallocation of the buffer use for promotion. It
 /// expects to get the same value that the default allocation method returned,
 /// i.e. result of a ViewOp.
-static LogicalResult deallocCallBack(OpBuilder &b, Value fullLocalView) {
+static LogicalResult
+defaultDeallocBufferCallBack(const LinalgPromotionOptions &options,
+                             OpBuilder &b, Value fullLocalView) {
   auto viewOp = fullLocalView.getDefiningOp<ViewOp>();
   assert(viewOp && "expected full local view to be a ViewOp");
-  std_dealloc(viewOp.source());
+  if (!options.useAlloca)
+    std_dealloc(viewOp.source());
   return success();
 }
 
@@ -182,11 +197,16 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
                             : [&](OpBuilder &builder, SubViewOp subViewOp,
                                   ArrayRef<Value> boundingSubViewSize,
                                   OperationFolder *folder) -> Optional<Value> {
-        return allocBufferCallBack(builder, subViewOp, boundingSubViewSize,
-                                   dynamicBuffers, alignment, folder);
+        return defaultAllocBufferCallBack(options, builder, subViewOp,
+                                          boundingSubViewSize, dynamicBuffers,
+                                          alignment, folder);
       });
   deallocationFn =
-      (options.deallocationFn ? *(options.deallocationFn) : deallocCallBack);
+      (options.deallocationFn
+           ? *(options.deallocationFn)
+           : [&](OpBuilder &b, Value buffer) {
+               return defaultDeallocBufferCallBack(options, b, buffer);
+             });
   auto defaultCopyCallBack = [&](OpBuilder &builder, Value src,
                                  Value dst) -> LogicalResult {
     linalg_copy(src, dst);
@@ -344,9 +364,8 @@ promoteSubViews(OpBuilder &b, LinalgOp op,
   }
 
   // 4. Dealloc all local buffers.
-  for (const auto &pi : *promotedBuffersAndViews) {
+  for (const auto &pi : *promotedBuffersAndViews)
     options.deallocationFn(b, pi.second.fullLocalView);
-  }
   return op;
 }
 
@@ -383,14 +402,17 @@ Optional<LinalgOp> mlir::linalg::promoteSubViews(OpBuilder &b,
 namespace {
 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
   LinalgPromotionPass() = default;
-  LinalgPromotionPass(bool dynamicBuffers) {
+  LinalgPromotionPass(bool dynamicBuffers, bool useAlloca) {
     this->dynamicBuffers = dynamicBuffers;
+    this->useAlloca = useAlloca;
   }
 
   void runOnFunction() override {
     OperationFolder folder(&getContext());
     getFunction().walk([this, &folder](LinalgOp op) {
-      auto options = LinalgPromotionOptions().setDynamicBuffers(dynamicBuffers);
+      auto options = LinalgPromotionOptions()
+                         .setDynamicBuffers(dynamicBuffers)
+                         .setUseAlloca(useAlloca);
       if (failed(promoteSubviewsPrecondition(op, options)))
         return;
       LLVM_DEBUG(llvm::dbgs() << "Promote: " << *(op.getOperation()) << "\n");
@@ -403,8 +425,8 @@ struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
 
 // TODO: support more transformation options in the pass.
 std::unique_ptr<OperationPass<FuncOp>>
-mlir::createLinalgPromotionPass(bool dynamicBuffers) {
-  return std::make_unique<LinalgPromotionPass>(dynamicBuffers);
+mlir::createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca) {
+  return std::make_unique<LinalgPromotionPass>(dynamicBuffers, useAlloca);
 }
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgPromotionPass() {
   return std::make_unique<LinalgPromotionPass>();

diff --git a/mlir/test/Dialect/Linalg/promote.mlir b/mlir/test/Dialect/Linalg/promote.mlir
index fe36755..ecf8e20 100644 (file)
--- a/mlir/test/Dialect/Linalg/promote.mlir
+++ b/mlir/test/Dialect/Linalg/promote.mlir
@@ -1,5 +1,6 @@
 // RUN: mlir-opt %s -linalg-promote-subviews | FileCheck %s
 // RUN: mlir-opt %s -linalg-promote-subviews="test-promote-dynamic" | FileCheck %s --check-prefix=DYNAMIC
+// RUN: mlir-opt %s -linalg-promote-subviews="test-use-alloca" | FileCheck %s --check-prefix=ALLOCA
 
 #map1 = affine_map<(d0) -> (d0 + 2)>
 #map2 = affine_map<(d0) -> (d0 + 4)>
@@ -45,16 +46,19 @@ func @matmul_f32(%A: memref<?xi8>, %M: index, %N: index, %K: index) {
 //       CHECK:         %[[vC:.*]] = subview {{.*}} : memref<?x?xf32>
 ///
 //       CHECK:         %[[tmpA:.*]] = alloc() : memref<32xi8>
+//      ALLOCA:         %[[tmpA:.*]] = alloca() : memref<32xi8>
 //       CHECK:         %[[fullA:.*]] = std.view %[[tmpA]][{{.*}}][{{.*}}] : memref<32xi8> to memref<?x?xf32>
 //     DYNAMIC:         std.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?x?xf32>
 //       CHECK:         %[[partialA:.*]] = subview %[[fullA]]{{.*}} : memref<?x?xf32> to memref<?x?xf32, #[[$strided2D_dynamic]]>
 ///
 //       CHECK:         %[[tmpB:.*]] = alloc() : memref<48xi8>
+//      ALLOCA:         %[[tmpB:.*]] = alloca() : memref<48xi8>
 //       CHECK:         %[[fullB:.*]] = std.view %[[tmpB]][{{.*}}][{{.*}}] : memref<48xi8> to memref<?x?xf32>
 //     DYNAMIC:         std.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?x?xf32>
 //       CHECK:         %[[partialB:.*]] = subview %[[fullB]]{{.*}} : memref<?x?xf32> to memref<?x?xf32, #[[$strided2D_dynamic]]>
 ///
 //       CHECK:         %[[tmpC:.*]] = alloc() : memref<24xi8>
+//      ALLOCA:         %[[tmpC:.*]] = alloca() : memref<24xi8>
 //       CHECK:         %[[fullC:.*]] = std.view %[[tmpC]][{{.*}}][{{.*}}] : memref<24xi8> to memref<?x?xf32>
 //     DYNAMIC:         std.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?x?xf32>
 //       CHECK:         %[[partialC:.*]] = subview %[[fullC]]{{.*}} : memref<?x?xf32> to memref<?x?xf32, #[[$strided2D_dynamic]]>
@@ -75,6 +79,9 @@ func @matmul_f32(%A: memref<?xi8>, %M: index, %N: index, %K: index) {
 //       CHECK:         dealloc %[[tmpA]] : memref<32xi8>
 //       CHECK:         dealloc %[[tmpB]] : memref<48xi8>
 //       CHECK:         dealloc %[[tmpC]] : memref<24xi8>
+//  ALLOCA-NOT:         dealloc %[[tmpA]] : memref<32xi8>
+//  ALLOCA-NOT:         dealloc %[[tmpB]] : memref<48xi8>
+//  ALLOCA-NOT:         dealloc %[[tmpC]] : memref<24xi8>
 
 // -----