[Mlir-commits] [mlir] [mlir][LinalgBlockPackMatmul] Add support for scalable block factors (PR #211354)
Stephen Long
llvmlistbot at llvm.org
Thu Jul 23 06:37:31 PDT 2026
https://github.com/steplong updated https://github.com/llvm/llvm-project/pull/211354
>From 3c8ac21e769cdb7712ca7317daa492651220bdb0 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at qti.qualcomm.com>
Date: Wed, 22 Jul 2026 11:04:40 -0700
Subject: [PATCH 1/4] [LinalgBlockPackMatmul] Add support for scalable block
factors
Assisted by: Claude Sonnet 4.6
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 8 +-
.../Dialect/Linalg/Transforms/Transforms.h | 5 +
.../Linalg/Transforms/BlockPackMatmul.cpp | 57 +++++++++++-
.../Linalg/block-pack-matmul-scalable.mlir | 91 +++++++++++++++++++
4 files changed, 155 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/block-pack-matmul-scalable.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 3a43af9ca1855..5573cf8b9feb7 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -215,10 +215,12 @@ def LinalgBlockPackMatmul : Pass<"linalg-block-pack-matmul"> {
%res = unpack %res_packed : 4D <MBxNBxmbxnb> -> 2D <MxN>
```
}];
- let dependentDialects = ["linalg::LinalgDialect", "tensor::TensorDialect"];
+ let dependentDialects = ["linalg::LinalgDialect", "tensor::TensorDialect",
+ "arith::ArithDialect",
+ "vector::VectorDialect"];
let options = [
- ListOption<"blockFactors", "block-factors", "int64_t",
- "Block factors (mb, nb, kb) for relayout">,
+ ListOption<"blockFactors", "block-factors", "std::string",
+ "Block factors (mb, nb, kb) for relayout; use [N] for scalable">,
Option<"allowPadding", "allow-padding", "bool",
/*default=*/"true",
"Allow packing padding">,
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 836682de4c404..9a59fb5ea7931 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1423,6 +1423,11 @@ struct BlockPackMatmulOptions {
/// the parallel dimensions and kb is the reduction dimension.
SmallVector<int64_t, 3> blockFactors;
+ /// Scalable flags for block factors. When true, the corresponding block
+ /// factor is multiplied by vector.vscale at runtime (SVE-style scalable).
+ /// Must be empty (all static) or have the same size as blockFactors.
+ SmallVector<bool, 3> scalableBlockFactors;
+
/// If true, allows packing of dimensions that only partially fit into the
/// block factors.
bool allowPadding = true;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
index 6ea1eb50b13ce..cfc91fe537ccd 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
@@ -8,9 +8,11 @@
#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/SmallVector.h"
@@ -161,8 +163,35 @@ linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
if (options->blockFactors.size() != 3)
return rewriter.notifyMatchFailure(linalgOp, "require 3 tile factors");
- SmallVector<OpFoldResult> mnkTiles =
- getAsOpFoldResult(rewriter.getI64ArrayAttr(options->blockFactors));
+ bool hasScalable = !options->scalableBlockFactors.empty();
+ if (hasScalable && options->scalableBlockFactors.size() != 3)
+ return rewriter.notifyMatchFailure(
+ linalgOp, "scalableBlockFactors must be empty or have 3 elements");
+
+ // Scalable tile sizes are non-constant at compile time, so they can never
+ // satisfy the full-tile divisibility check. Reject early before creating
+ // any ops to avoid modifying IR before returning notifyMatchFailure.
+ if (!options->allowPadding && hasScalable)
+ return rewriter.notifyMatchFailure(
+ linalgOp, "scalable block factors require allow-padding=true");
+
+ // Build OpFoldResult tile sizes. Scalable dimensions are emitted as
+ // arith.constant N : index multiplied by vector.vscale.
+ SmallVector<OpFoldResult> mnkTiles;
+ for (auto [i, factor] : llvm::enumerate(options->blockFactors)) {
+ bool isScalable = hasScalable && options->scalableBlockFactors[i];
+ if (!isScalable) {
+ mnkTiles.push_back(rewriter.getIndexAttr(factor));
+ continue;
+ }
+ Value cst = arith::ConstantIndexOp::create(rewriter, linalgOp.getLoc(),
+ factor);
+ Value vscale = vector::VectorScaleOp::create(rewriter, linalgOp.getLoc(),
+ rewriter.getIndexType());
+ mnkTiles.push_back(
+ arith::MulIOp::create(rewriter, linalgOp.getLoc(), cst, vscale)
+ .getResult());
+ }
// If padding is disabled, make sure that dimensions can be packed cleanly.
if (!options->allowPadding &&
@@ -300,7 +329,29 @@ struct LinalgBlockPackMatmul
ControlBlockPackMatmulFn controlFn =
[&](linalg::LinalgOp op) -> BlockPackMatmulOptions {
BlockPackMatmulOptions options;
- options.blockFactors = SmallVector<int64_t>{*blockFactors};
+
+ // Parse block-factors strings. Each element is either "N" (static) or
+ // "[N]" (scalable, i.e. N * vscale at runtime).
+ for (const std::string &f : *blockFactors) {
+ StringRef s(f);
+ if (s.starts_with("[") && s.ends_with("]")) {
+ int64_t val = 0;
+ s.drop_front().drop_back().getAsInteger(10, val);
+ options.blockFactors.push_back(val);
+ options.scalableBlockFactors.push_back(true);
+ } else {
+ int64_t val = 0;
+ s.getAsInteger(10, val);
+ options.blockFactors.push_back(val);
+ options.scalableBlockFactors.push_back(false);
+ }
+ }
+ // If all flags are false, clear the vector so blockPackMatmul can take
+ // the cheaper static path.
+ if (llvm::none_of(options.scalableBlockFactors,
+ [](bool b) { return b; }))
+ options.scalableBlockFactors.clear();
+
options.allowPadding = allowPadding;
options.mnkPaddedSizesNextMultipleOf =
SmallVector<int64_t>{*mnkPaddedSizesNextMultipleOf};
diff --git a/mlir/test/Dialect/Linalg/block-pack-matmul-scalable.mlir b/mlir/test/Dialect/Linalg/block-pack-matmul-scalable.mlir
new file mode 100644
index 0000000000000..a5829c210d1a8
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/block-pack-matmul-scalable.mlir
@@ -0,0 +1,91 @@
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=[32],[16],[64] allow-padding=1" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=SCALABLE
+
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=[32],[16],[64] allow-padding=0" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=SCALABLE-NOPAD
+
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=32,16,[64] allow-padding=1" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=MIXED
+
+// -----
+
+// All-scalable block factors, static input shapes, allow-padding=1.
+// Inner tile sizes and outer dimensions are all dynamic (vscale-relative).
+
+func.func @block_matmul_scalable_static(
+ %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {
+ %0 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)
+ outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>
+ return %0 : tensor<128x128xf32>
+}
+
+// SCALABLE-LABEL: func @block_matmul_scalable_static(
+// SCALABLE-DAG: %[[C32:.+]] = arith.constant 32 : index
+// SCALABLE-DAG: %[[C16:.+]] = arith.constant 16 : index
+// SCALABLE-DAG: %[[C64:.+]] = arith.constant 64 : index
+// SCALABLE-DAG: %[[MB:.+]] = arith.muli %{{.+}}, %[[C32]] : index
+// SCALABLE-DAG: %[[NB:.+]] = arith.muli %{{.+}}, %[[C16]] : index
+// SCALABLE-DAG: %[[KB:.+]] = arith.muli %{{.+}}, %[[C64]] : index
+// SCALABLE: linalg.pack %{{.+}} outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[KB]]]
+// SCALABLE: linalg.pack %{{.+}} outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [%[[NB]], %[[KB]]]
+// SCALABLE: linalg.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[NB]]]
+// SCALABLE: linalg.generic
+// SCALABLE: linalg.unpack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[NB]]]
+
+// Scalable factors with allow-padding=0: transform does not apply.
+// SCALABLE-NOPAD-LABEL: func @block_matmul_scalable_static(
+// SCALABLE-NOPAD-NOT: linalg.pack
+// SCALABLE-NOPAD: linalg.matmul
+// SCALABLE-NOPAD-NOT: linalg.unpack
+
+// -----
+
+// All-scalable block factors, dynamic input shapes, allow-padding=1.
+// Both outer tile counts and inner tile sizes are fully dynamic.
+
+func.func @block_matmul_scalable_dynamic(
+ %A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%C : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// SCALABLE-LABEL: func @block_matmul_scalable_dynamic(
+// SCALABLE-DAG: %[[C32:.+]] = arith.constant 32 : index
+// SCALABLE-DAG: %[[C16:.+]] = arith.constant 16 : index
+// SCALABLE-DAG: %[[C64:.+]] = arith.constant 64 : index
+// SCALABLE-DAG: %[[MB:.+]] = arith.muli %{{.+}}, %[[C32]] : index
+// SCALABLE-DAG: %[[NB:.+]] = arith.muli %{{.+}}, %[[C16]] : index
+// SCALABLE-DAG: %[[KB:.+]] = arith.muli %{{.+}}, %[[C64]] : index
+// SCALABLE: linalg.pack %{{.+}} outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[KB]]]
+// SCALABLE: linalg.pack %{{.+}} outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [%[[NB]], %[[KB]]]
+// SCALABLE: linalg.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[NB]]]
+// SCALABLE: linalg.generic
+// SCALABLE: linalg.unpack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [%[[MB]], %[[NB]]]
+
+// SCALABLE-NOPAD-LABEL: func @block_matmul_scalable_dynamic(
+// SCALABLE-NOPAD-NOT: linalg.pack
+// SCALABLE-NOPAD: linalg.matmul
+// SCALABLE-NOPAD-NOT: linalg.unpack
+
+// -----
+
+// Mixed block factors (mb=32, nb=16 static; kb=[64] scalable), static input.
+// Only the K-dimension inner tile is scalable; mb and nb remain static.
+
+func.func @block_matmul_mixed_static(
+ %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {
+ %0 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)
+ outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>
+ return %0 : tensor<128x128xf32>
+}
+
+// MIXED-LABEL: func @block_matmul_mixed_static(
+// MIXED-DAG: %[[C64:.+]] = arith.constant 64 : index
+// MIXED-DAG: %[[KB:.+]] = arith.muli %{{.+}}, %[[C64]] : index
+// MIXED: linalg.pack %{{.+}} outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [32, %[[KB]]]
+// MIXED: linalg.pack %{{.+}} outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, %[[KB]]]
+// MIXED: linalg.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [32, 16]
+// MIXED: linalg.generic
+// MIXED: linalg.unpack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [32, 16]
+
>From 46852e32748c6ffab6973f12de9c06e5c6d6d060 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at qti.qualcomm.com>
Date: Wed, 22 Jul 2026 14:34:41 -0700
Subject: [PATCH 2/4] clang-format PR
---
mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
index cfc91fe537ccd..ebc287463bec2 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
@@ -184,8 +184,8 @@ linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
mnkTiles.push_back(rewriter.getIndexAttr(factor));
continue;
}
- Value cst = arith::ConstantIndexOp::create(rewriter, linalgOp.getLoc(),
- factor);
+ Value cst =
+ arith::ConstantIndexOp::create(rewriter, linalgOp.getLoc(), factor);
Value vscale = vector::VectorScaleOp::create(rewriter, linalgOp.getLoc(),
rewriter.getIndexType());
mnkTiles.push_back(
@@ -348,8 +348,7 @@ struct LinalgBlockPackMatmul
}
// If all flags are false, clear the vector so blockPackMatmul can take
// the cheaper static path.
- if (llvm::none_of(options.scalableBlockFactors,
- [](bool b) { return b; }))
+ if (llvm::none_of(options.scalableBlockFactors, [](bool b) { return b; }))
options.scalableBlockFactors.clear();
options.allowPadding = allowPadding;
>From c80811912e002c40c448ab690c13a6b97a5f5b4d Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Thu, 23 Jul 2026 09:37:08 -0400
Subject: [PATCH 3/4] Update
mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
Co-authored-by: Adam Siemieniuk <adam.siemieniuk at intel.com>
---
mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
index ebc287463bec2..e6147c8814b9d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
@@ -173,7 +173,7 @@ linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
// any ops to avoid modifying IR before returning notifyMatchFailure.
if (!options->allowPadding && hasScalable)
return rewriter.notifyMatchFailure(
- linalgOp, "scalable block factors require allow-padding=true");
+ linalgOp, "scalable block factors require padding");
// Build OpFoldResult tile sizes. Scalable dimensions are emitted as
// arith.constant N : index multiplied by vector.vscale.
>From 3f966c039d99578847d20ae6d06318c6f4ef3bb7 Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Thu, 23 Jul 2026 09:37:21 -0400
Subject: [PATCH 4/4] Update
mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
Co-authored-by: Adam Siemieniuk <adam.siemieniuk at intel.com>
---
mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
index e6147c8814b9d..ba64b6b2e4447 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BlockPackMatmul.cpp
@@ -178,7 +178,7 @@ linalg::blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
// Build OpFoldResult tile sizes. Scalable dimensions are emitted as
// arith.constant N : index multiplied by vector.vscale.
SmallVector<OpFoldResult> mnkTiles;
- for (auto [i, factor] : llvm::enumerate(options->blockFactors)) {
+ for (auto [idx, factor] : llvm::enumerate(options->blockFactors)) {
bool isScalable = hasScalable && options->scalableBlockFactors[i];
if (!isScalable) {
mnkTiles.push_back(rewriter.getIndexAttr(factor));
More information about the Mlir-commits
mailing list