[mlir][linalg] Allow promotion to use the original subview size #144334

zbenzion · 2025-06-16T11:37:51Z

linalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic.

Adding a promotion option to use the original subview size by default, thus minimizing the allocation size.
Fixes #144268.

github-actions · 2025-06-16T11:38:12Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2025-06-16T11:38:41Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: None (zbenzion)

Changes

linalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic.

Adding a promotion option to use the original subview size by default, thus minimizing the allocation size.

Full diff: https://github.com/llvm/llvm-project/pull/144334.diff

5 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+1)
(modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+10-1)
(modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+3)
(modified) mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp (+10-3)
(modified) mlir/test/Dialect/Linalg/promotion_options.mlir (+56)

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 15ea5e7bf7159..42de4a2f153c9 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1309,6 +1309,7 @@ def PromoteOp : Op<Transform_Dialect, "structured.promote",
                        DefaultValuedAttr<I64ArrayAttr, "{}">:$operands_to_promote,
                        DefaultValuedAttr<BoolArrayAttr, "{}">:$use_full_tile_buffers,
                        UnitAttr:$use_full_tiles_by_default,
+                       UnitAttr:$use_original_subview_size,
                        UnitAttr:$use_alloca,
                        OptionalAttr<AnyAttr>:$memory_space,
                        OptionalAttr<DeviceMappingArrayAttr>:$mapping,
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 2eef0a06d0eb4..d9dceade0ac95 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -381,6 +381,13 @@ struct LinalgPromotionOptions {
     useFullTileBuffersDefault = use;
     return *this;
   }
+  /// If true bufffers will be allocated with the original subview size. This
+  /// may result in more dynamic allocations, in case of dynamic sizes.
+  bool useOriginalSubviewSize = false;
+  LinalgPromotionOptions &setUseOriginalSubviewSize(bool originalSize) {
+    useOriginalSubviewSize = originalSize;
+    return *this;
+  }
   /// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
   std::optional<unsigned> alignment;
   LinalgPromotionOptions &setAlignment(unsigned align) {
@@ -712,7 +719,8 @@ FailureOr<LinalgOp> specializeGenericOp(RewriterBase &rewriter,
                                         GenericOp genericOp);
 
 /// Create a new buffer using the `allocationFn` provided. The size of this
-/// buffer is the smallest constant bounding size along each dimension that
+/// buffer is either the original subview size when 'useOriginalSubviewSize' is
+/// set to true or the smallest constant bounding size along each dimension that
 /// can be computed for the size of the result of `subView`. Returns the
 /// allocated buffer as `fullLocalView` and the view that matches the size of
 /// the result of subview operation as `partialLocalView`.
@@ -722,6 +730,7 @@ struct PromotionInfo {
 };
 FailureOr<PromotionInfo>
 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
+                          bool useOriginalSubviewSize,
                           const AllocBufferCallbackFn &allocationFn,
                           DataLayout &layout);
 
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index b2c28f5eed33c..6be73a11e1b54 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2236,6 +2236,9 @@ transform::PromoteOp::applyToOne(transform::TransformRewriter &rewriter,
   if (getUseFullTilesByDefault())
     promotionOptions = promotionOptions.setUseFullTileBuffersByDefault(
         getUseFullTilesByDefault());
+  if (getUseOriginalSubviewSize())
+    promotionOptions =
+        promotionOptions.setUseOriginalSubviewSize(getUseOriginalSubviewSize());
   if (getUseAlloca())
     promotionOptions = promotionOptions.setUseAlloca(getUseAlloca());
   if (!getUseFullTileBuffers().empty())
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index d599ddd220dde..80b37544a99f8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -148,6 +148,9 @@ struct LinalgOpInstancePromotionOptions {
   llvm::SmallSet<int64_t, 4> operandsNumbersToCopyIn;
   /// True if the full view should be used for the promoted buffer.
   DenseMap<Value, bool> useFullTileBuffers;
+  /// True if the original subview size should be used. This means the full tile
+  /// buffer is the same size as the partial view.
+  bool useOriginalSubviewSize;
 
   /// Callback functions for allocation and deallocation of promoted buffers, as
   /// well as to copy the data into and out of these buffers.
@@ -170,6 +173,7 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
       options.useFullTileBuffers.value_or(llvm::SmallBitVector());
   vUseFullTileBuffers.resize(linalgOp->getNumOperands(),
                              options.useFullTileBuffersDefault);
+  useOriginalSubviewSize = options.useOriginalSubviewSize;
 
   for (OpOperand &opOperand : linalgOp->getOpOperands()) {
     int64_t operandNumber = opOperand.getOperandNumber();
@@ -237,7 +241,8 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
 // by a partial `copy` op.
 FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
     OpBuilder &b, Location loc, memref::SubViewOp subView,
-    const AllocBufferCallbackFn &allocationFn, DataLayout &layout) {
+    bool useOriginalSubviewSize, const AllocBufferCallbackFn &allocationFn,
+    DataLayout &layout) {
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
   SmallVector<Value, 4> fullSizes;
@@ -254,7 +259,8 @@ FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
     // to look for the bound.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
     Value size;
-    if (auto attr = llvm::dyn_cast_if_present<Attribute>(rangeValue.size)) {
+    if (llvm::isa_and_present<Attribute>(rangeValue.size) ||
+        useOriginalSubviewSize) {
       size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
     } else {
       FailureOr<int64_t> upperBound =
@@ -295,7 +301,8 @@ promoteSubViews(ImplicitLocOpBuilder &b,
     memref::SubViewOp subView =
         cast<memref::SubViewOp>(v.second.getDefiningOp());
     auto promotionInfo = promoteSubviewAsNewBuffer(
-        b, b.getLoc(), subView, options.allocationFn, layout);
+        b, b.getLoc(), subView, options.useOriginalSubviewSize,
+        options.allocationFn, layout);
     if (failed(promotionInfo))
       return failure();
     promotionInfoMap[v.first] = *promotionInfo;
diff --git a/mlir/test/Dialect/Linalg/promotion_options.mlir b/mlir/test/Dialect/Linalg/promotion_options.mlir
index caa72ba24316f..dbc073c2665f9 100644
--- a/mlir/test/Dialect/Linalg/promotion_options.mlir
+++ b/mlir/test/Dialect/Linalg/promotion_options.mlir
@@ -42,3 +42,59 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+func.func @matmul_f32(%A: memref<512x256xf32>, %B: memref<256x512xf32>, %C: memref<256x256xf32>, %s0: index, %s1: index, %s2: index) {
+  %c0 = arith.constant 0 : index
+  %c256 = arith.constant 256 : index
+  %c512 = arith.constant 512 : index
+  scf.for %arg4 = %c0 to %c512 step %s0 {
+    scf.for %arg5 = %c0 to %c512 step %s1 {
+      scf.for %arg6 = %c0 to %c256 step %s2 {
+        %i0 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg4)[%s0]
+        %i1 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg5)[%s1]
+        %i2 = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%arg6)[%s2]
+        %0 = memref.subview %A[%arg4, %arg6][%i0, %i2][1, 1] : memref<512x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>
+        %1 = memref.subview %B[%arg6, %arg5][%i2, %i1][1, 1] : memref<256x512xf32> to memref<?x?xf32, strided<[512, 1], offset: ?>>
+        %2 = memref.subview %C[%arg4, %arg5][%i0, %i1][1, 1] : memref<256x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>
+        linalg.matmul
+          ins(%0, %1: memref<?x?xf32, strided<[256, 1], offset: ?>>,
+                      memref<?x?xf32, strided<[512, 1], offset: ?>>)
+          outs(%2: memref<?x?xf32, strided<[256, 1], offset: ?>>)
+      }
+    }
+  }
+  return
+}
+
+// CHECK-LABEL:   func.func @matmul_f32(
+// CHECK-SAME:      %[[ARG0:.*]]: memref<512x256xf32>
+// CHECK-SAME:      %[[ARG1:.*]]: memref<256x512xf32>
+// CHECK-SAME:      %[[ARG2:.*]]: memref<256x256xf32>
+// CHECK-SAME:      %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index
+// CHECK:           %[[C4:.*]] = arith.constant 4 : index
+
+// CHECK:                 %[[i0:.*]] = affine.min
+// CHECK:                 %[[i1:.*]] = affine.min
+// CHECK:                 %[[i2:.*]] = affine.min
+
+// CHECK:                 %[[VAL_13:.*]] = arith.muli %[[i0]], %[[i2]] : index
+// CHECK:                 %[[VAL_14:.*]] = arith.muli %[[VAL_13]], %[[C4]] : index
+// CHECK:                 %[[VAL_15:.*]] = memref.alloc(%[[VAL_14]]) : memref<?xi8>
+
+// CHECK:                 %[[VAL_18:.*]] = arith.muli %[[i2]], %[[i1]] : index
+// CHECK:                 %[[VAL_19:.*]] = arith.muli %[[VAL_18]], %[[C4]] : index
+// CHECK:                 %[[VAL_20:.*]] = memref.alloc(%[[VAL_19]]) : memref<?xi8>
+
+// CHECK:                 %[[VAL_23:.*]] = arith.muli %[[i0]], %[[i1]] : index
+// CHECK:                 %[[VAL_24:.*]] = arith.muli %[[VAL_23]], %[[C4]] : index
+// CHECK:                 %[[VAL_25:.*]] = memref.alloc(%[[VAL_24]]) : memref<?xi8>
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1 = transform.structured.promote %0 { use_original_subview_size } : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}

ftynse

LGTM after comments addressed, thanks!

ftynse · 2025-06-24T07:02:59Z

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

@@ -427,6 +427,13 @@ struct LinalgPromotionOptions {
    useFullTileBuffersDefault = use;
    return *this;
  }
+  /// If true bufffers will be allocated with the original subview size. This


Suggested change

/// If true bufffers will be allocated with the original subview size. This

/// If true, buffers will be allocated with the original subview size. This

ftynse · 2025-06-24T07:04:14Z

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

@@ -811,6 +819,7 @@ struct PromotionInfo {
 };
 FailureOr<PromotionInfo>
 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
+                          bool useOriginalSubviewSize,


Let's have this take const LinalgPromotionOptions & instead of an individual boolean.

In general, sounds good!

It seems though that the caller promoteSubViews is already using LinalgOpInstancePromotionOptions, which is an internal representation in Promotion.cpp.
Another issue is that LinalgPromotionOptions doesn't force an explicit allocationFn, and promoteSubviewAsNewBuffer does.

We could maybe, alternatively, provide a LinalgPromoteSubviewAsNewBufferOptions, which will pack the boolean and the allocation function.
What do you think?

Can we just extend one or another? I'd rather avoid yet another helper class so if it's impossible to extend one of the existing classes, we can stick with the single argument here.

Since this method is used both internally in Promotion.cpp (LinalgOpInstancePromotionOptions) and is exposed externally (LinalgPromotionOptions) I feel like it'd be a little unnatural to use one or the other at this stage. I kept the single argument, as you suggested.

If everything looks good, I'd appreciate if you can hit the merge button for me, since I don't have write access.
Thanks for the review!

linalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic. Adding a promotion option to use the original subview size by default, thus minimizing the allocation size.

zbenzion requested review from ftynse, nicolasvasilache, dcaballe and rengolin as code owners June 16, 2025 11:37

llvmbot added mlir:linalg mlir labels Jun 16, 2025

rengolin requested a review from matthias-springer June 16, 2025 11:39

zbenzion mentioned this pull request Jun 19, 2025

[mlir][linalg][promotion] Allow promoting dynamic buffer with dynamic alloc size #144268

Open

zbenzion force-pushed the linalg/promote-orig-size branch from e990a31 to 629ff4a Compare June 23, 2025 07:09

ftynse approved these changes Jun 24, 2025

View reviewed changes

zbenzion force-pushed the linalg/promote-orig-size branch from 629ff4a to 981ee07 Compare June 29, 2025 07:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][linalg] Allow promotion to use the original subview size #144334

[mlir][linalg] Allow promotion to use the original subview size #144334

zbenzion commented Jun 16, 2025 •

edited by CoTinker

Loading

Uh oh!

github-actions bot commented Jun 16, 2025

Uh oh!

llvmbot commented Jun 16, 2025 •

edited

Loading

Uh oh!

ftynse left a comment

Uh oh!

ftynse Jun 24, 2025

Uh oh!

zbenzion Jun 29, 2025

Uh oh!

ftynse Jun 24, 2025

Uh oh!

zbenzion Jun 24, 2025

Uh oh!

ftynse Jun 27, 2025

Uh oh!

zbenzion Jun 29, 2025

Uh oh!

Uh oh!

	/// If true bufffers will be allocated with the original subview size. This
	/// If true, buffers will be allocated with the original subview size. This

[mlir][linalg] Allow promotion to use the original subview size #144334

Are you sure you want to change the base?

[mlir][linalg] Allow promotion to use the original subview size #144334

Conversation

zbenzion commented Jun 16, 2025 • edited by CoTinker Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jun 16, 2025

Uh oh!

llvmbot commented Jun 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

ftynse Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

zbenzion Jun 29, 2025

Choose a reason for hiding this comment

Uh oh!

ftynse Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

zbenzion Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

ftynse Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

zbenzion Jun 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zbenzion commented Jun 16, 2025 •

edited by CoTinker

Loading

llvmbot commented Jun 16, 2025 •

edited

Loading