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

[nicolasvasilache]

Summary

A relevant test is also added.

[bkramer]

Thanks

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 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
 /// std.load/std.store accesses.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToLoopsPass();
 
 /// Create a pass to convert Linalg operations to scf.parallel loops and
 /// std.load/std.store accesses.

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgLowerToParallelLoops
   let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";
 }
 
 def LinalgPromotion : FunctionPass<"linalg-promote-subviews"> {
   let summary = "Promote subview ops to local buffers";
   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.">
   ];
 }
 
 def LinalgTiling : FunctionPass<"linalg-tile"> {
   let summary = "Tile operations in the linalg dialect";
   let constructor = "mlir::createLinalgTilingPass()";
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t",

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

LinalgPromotionOptions &setDynamicBuffers(unsigned dynamic) {
     return *this;
   }
   /// Alignment of promoted buffer. If `None` do not specify alignment.
   Optional<unsigned> alignment = None;
   LinalgPromotionOptions &setAlignment(unsigned align) {
     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.
   Optional<AllocBufferCallbackFn> allocationFn = None;
   Optional<DeallocBufferCallbackFn> deallocationFn = None;
   LinalgPromotionOptions &
   setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn,
                                DeallocBufferCallbackFn const &deallocFn) {
     allocationFn = allocFn;
     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;
   Optional<CopyCallbackFn> copyOutFn = None;
   LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
                                           CopyCallbackFn const &copyOut) {
     copyInFn = copyIn;
     copyOutFn = copyOut;

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc,
   }
   return boundingConst && *boundingConst >= 0
              ? b.create<ConstantIndexOp>(loc, *boundingConst)
              : size;
 }
 
 /// 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();
   auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
   IntegerAttr alignment_attr;
   if (alignment.hasValue())
     alignment_attr =
         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,
   ShapedType viewType = subView.getType();
   int64_t rank = viewType.getRank();
   (void)rank;
   assert(rank > 0 && boundingSubViewSize.size() == static_cast<size_t>(rank));
   auto zero = folded_std_constant_index(folder, 0);
   auto one = folded_std_constant_index(folder, 1);
 
   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);
   Value view = folded_std_view(
       folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer,
       zero, boundingSubViewSize);
   return view;
 }
 
 /// 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");
+  if (!options.useAlloca)
     std_dealloc(viewOp.source());
   return success();
 }
 
 namespace {
 
 /// Helper struct that captures the information required to apply the
 /// transformation on each op. This bridges the abstraction gap with the
 /// user-facing API which exposes positional arguments to control which operands
for (unsigned idx = 0; idx != nBuffers; ++idx) {
     }
   }
 
   allocationFn =
       (options.allocationFn ? *(options.allocationFn)
                             : [&](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);
     return success();
   };
   copyInFn = (options.copyInFn ? *(options.copyInFn) : defaultCopyCallBack);
   copyOutFn = (options.copyOutFn ? *(options.copyOutFn) : defaultCopyCallBack);
 }
promoteSubViews(OpBuilder &b, LinalgOp op,
   // 3. Emit write-back for the promoted output views: copy the partial view.
   for (auto viewAndPartialLocalView : writebackViews) {
     if (failed(options.copyOutFn(b, viewAndPartialLocalView.second,
                                  viewAndPartialLocalView.first)))
       return {};
   }
 
   // 4. Dealloc all local buffers.
-  for (const auto &pi : *promotedBuffersAndViews) {
+  for (const auto &pi : *promotedBuffersAndViews)
     options.deallocationFn(b, pi.second.fullLocalView);
-  }
   return op;
 }
 
 LogicalResult
 mlir::linalg::promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options) {
   LinalgOp linOp = dyn_cast<LinalgOp>(op);
   // Transformation applies to buffers only.
Optional<LinalgOp> mlir::linalg::promoteSubViews(OpBuilder &b,
   LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options);
   return ::promoteSubViews(
       b, linalgOp, LinalgOpInstancePromotionOptions(linalgOp, options), folder);
 }
 
 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");
       OpBuilder b(op);
       promoteSubViews(b, op, options, &folder);
     });
   }
 };
 } // namespace
 
 // 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>();
 }

mlir/test/Dialect/Linalg/promote.mlir

 // 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)>
 #map3 = affine_map<(d0) -> (d0 + 3)>
 
 // CHECK-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[$strided2D_dynamic:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 

 //       CHECK:   scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //       CHECK:     scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //       CHECK:       scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //       CHECK:         %[[vA:.*]] = subview {{.*}} : memref<?x?xf32>
 //       CHECK:         %[[vB:.*]] = subview {{.*}} : memref<?x?xf32>
 //       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]]>
 
 //       CHECK:         linalg.copy(%[[vA]], %[[partialA]]) : memref<?x?xf32, #[[$strided2D]]>, memref<?x?xf32, #[[$strided2D_dynamic]]>
 //       CHECK:         linalg.copy(%[[vB]], %[[partialB]]) : memref<?x?xf32, #[[$strided2D]]>, memref<?x?xf32, #[[$strided2D_dynamic]]>
 //       CHECK:         linalg.copy(%[[vC]], %[[partialC]]) : memref<?x?xf32, #[[$strided2D]]>, memref<?x?xf32, #[[$strided2D_dynamic]]>
 //
 //       CHECK:         linalg.matmul %[[partialA]], %[[partialB]], %[[partialC]] :
 //       CHECK:           memref<?x?xf32, #[[$strided2D_dynamic]]>,
 //       CHECK:           memref<?x?xf32, #[[$strided2D_dynamic]]>,
 //       CHECK:           memref<?x?xf32, #[[$strided2D_dynamic]]>
 //
 //       CHECK:         linalg.copy(%[[partialC]], %[[vC]]) :
 //       CHECK:           memref<?x?xf32, #[[$strided2D_dynamic]]>,
 //       CHECK:           memref<?x?xf32, #[[$strided2D]]>
 //
 //       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>
 
 // -----
 
 func @matmul_f64(%A: memref<?xi8>, %M: index, %N: index, %K: index) {
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %c0 = constant 0 : index