[Mlir-commits] [mlir] [mlir][xegpu] Support batched matmul in VectorToXeGPU ContractionLowering (PR #211947)
Jianhui Li
llvmlistbot at llvm.org
Fri Jul 24 19:51:03 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/211947
>From d93be98f0671d505e7c28a55aefcc1b82ae5b2a9 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 24 Jul 2026 22:20:54 +0000
Subject: [PATCH 1/2] [mlir][xegpu] Support batched matmul in VectorToXeGPU
ContractionLowering
Generalize the row-major matmul check in ContractionLowering so that
(batched) N-D vector.contract ops lower to xegpu.dpas, where any leading
dimensions are batch dimensions shared by lhs, rhs, and acc and the
innermost two dimensions are the core matmul dims (M, K), (K, N), (M, N).
Contractions that do not match this pattern (e.g. an N-D contract whose
leading dim is not a shared batch dim) get a soft match failure with a
descriptive reason instead of silently hitting the old "Expects lhs and
rhs 2D vectors" message.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 81 ++++++++++++++++---
.../VectorToXeGPU/contract-to-xegpu.mlir | 48 +++++++++++
2 files changed, 118 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 727c98aef6619..850fd80793736 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -924,6 +924,63 @@ struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
}
};
+// If `indexingMaps` describe a (batched) row-major matmul
+// lhs[b..., m, k], rhs[b..., k, n], acc[b..., m, n]
+// return the number of leading batch dims (0 for a plain 2D matmul);
+// otherwise return std::nullopt.
+static std::optional<int64_t>
+getRowMajorMatmulBatchRank(ArrayAttr indexingMaps) {
+ if (indexingMaps.size() != 3)
+ return std::nullopt;
+
+ AffineMap mapA = cast<AffineMapAttr>(indexingMaps[0]).getValue();
+ AffineMap mapB = cast<AffineMapAttr>(indexingMaps[1]).getValue();
+ AffineMap mapC = cast<AffineMapAttr>(indexingMaps[2]).getValue();
+
+ // The result map exposes the batch dims followed by the core (m, n) dims.
+ if (mapC.getNumResults() < 2)
+ return std::nullopt;
+ int64_t batchRank = mapC.getNumResults() - 2;
+
+ // A single `k` reduction gives batchRank + 3 iteration dims; each operand
+ // map exposes batchRank + 2 dims (batch dims + 2 core dims).
+ unsigned numDims = static_cast<unsigned>(batchRank) + 3;
+ unsigned numOperandResults = static_cast<unsigned>(batchRank) + 2;
+ if (mapA.getNumInputs() != numDims || mapB.getNumInputs() != numDims ||
+ mapC.getNumInputs() != numDims)
+ return std::nullopt;
+ if (mapA.getNumResults() != numOperandResults ||
+ mapB.getNumResults() != numOperandResults)
+ return std::nullopt;
+
+ // Reconstruct the canonical maps from the batch/m/n dims of the result and
+ // the k dim of lhs, then compare against the actual maps.
+ MLIRContext *context = indexingMaps.getContext();
+ ArrayRef<AffineExpr> batchDims = mapC.getResults().take_front(batchRank);
+ AffineExpr m = mapC.getResult(batchRank);
+ AffineExpr n = mapC.getResult(batchRank + 1);
+ AffineExpr k = mapA.getResult(batchRank + 1);
+
+ SmallVector<AffineExpr> aDims = llvm::to_vector(batchDims);
+ aDims.push_back(m);
+ aDims.push_back(k);
+ SmallVector<AffineExpr> bDims = llvm::to_vector(batchDims);
+ bDims.push_back(k);
+ bDims.push_back(n);
+ SmallVector<AffineExpr> cDims = llvm::to_vector(batchDims);
+ cDims.push_back(m);
+ cDims.push_back(n);
+
+ auto expected = ArrayAttr::get(
+ context,
+ {AffineMapAttr::get(AffineMap::get(numDims, 0, aDims, context)),
+ AffineMapAttr::get(AffineMap::get(numDims, 0, bDims, context)),
+ AffineMapAttr::get(AffineMap::get(numDims, 0, cDims, context))});
+ if (indexingMaps != expected)
+ return std::nullopt;
+ return batchRank;
+}
+
struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
using Base::Base;
@@ -935,21 +992,23 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
return rewriter.notifyMatchFailure(contractOp,
"Expects add combining kind");
+ TypedValue<VectorType> lhs = contractOp.getLhs();
+ TypedValue<VectorType> rhs = contractOp.getRhs();
TypedValue<Type> acc = contractOp.getAcc();
VectorType accType = dyn_cast<VectorType>(acc.getType());
- if (!accType || accType.getRank() != 2)
- return rewriter.notifyMatchFailure(contractOp, "Expects acc 2D vector");
+ if (!accType)
+ return rewriter.notifyMatchFailure(contractOp, "Expects vector acc");
- // Accept only plain 2D data layout.
- // VNNI packing is applied to DPAS as a separate lowering step.
- TypedValue<VectorType> lhs = contractOp.getLhs();
- TypedValue<VectorType> rhs = contractOp.getRhs();
- if (lhs.getType().getRank() != 2 || rhs.getType().getRank() != 2)
- return rewriter.notifyMatchFailure(contractOp,
- "Expects lhs and rhs 2D vectors");
+ if (!getRowMajorMatmulBatchRank(contractOp.getIndexingMapsAttr()))
+ return rewriter.notifyMatchFailure(
+ contractOp, "Expects a (batched) row-major matmul: leading dims must "
+ "be batch dims shared by lhs, rhs, and acc; innermost two "
+ "dims must be (M, K), (K, N), and (M, N)");
- if (!isRowMajorMatmul(contractOp.getIndexingMapsAttr()))
- return rewriter.notifyMatchFailure(contractOp, "Invalid indexing maps");
+ // xegpu.dpas operands are limited to rank 4 (2 batch + 2 core dims).
+ if (accType.getRank() > 4)
+ return rewriter.notifyMatchFailure(contractOp,
+ "Expects operands of rank 4 or less");
auto dpasOp = xegpu::DpasOp::create(rewriter, loc,
TypeRange{contractOp.getResultType()},
diff --git a/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
index 292e4ff882000..7a5bc8f95f3c7 100644
--- a/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/contract-to-xegpu.mlir
@@ -180,3 +180,51 @@ func.func @negative_accumulator_shape(%lhs: vector<8x16xf16>, %rhs: vector<16x16
// CHECK-LABEL: @negative_accumulator_shape(
// CHECK: vector.contract
+
+// -----
+
+// A batched matmul whose leading dimension is a batch dimension shared across
+// lhs, rhs, and acc lowers to a batched xegpu.dpas.
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+func.func @dpas_batched_gemm(%lhs: vector<4x8x16xf16>, %rhs: vector<4x16x16xf16>,
+ %acc: vector<4x8x16xf32>) -> vector<4x8x16xf32> {
+ %3 = vector.contract
+ {indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>} %lhs, %rhs, %acc
+ : vector<4x8x16xf16>, vector<4x16x16xf16> into vector<4x8x16xf32>
+ return %3 : vector<4x8x16xf32>
+}
+
+// CHECK-LABEL: @dpas_batched_gemm(
+// CHECK-SAME: %[[LHS:.+]]: vector<4x8x16xf16>,
+// CHECK-SAME: %[[RHS:.+]]: vector<4x16x16xf16>,
+// CHECK-SAME: %[[ACC:.+]]: vector<4x8x16xf32>
+// CHECK: %[[DPAS:.+]] = xegpu.dpas
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]]
+// CHECK-SAME: {{.*}}-> vector<4x8x16xf32>
+// CHECK: return %[[DPAS]]
+
+// -----
+
+// An N-D contraction whose leading dimension is not a batch dimension shared
+// across lhs, rhs, and acc does not map to a (batched) row-major matmul.
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+func.func @negative_non_batched_nd(%lhs: vector<128x64x64xf16>, %rhs: vector<64x64xf16>,
+ %acc: vector<128x64xf16>) -> vector<128x64xf16> {
+ %3 = vector.contract
+ {indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>} %lhs, %rhs, %acc
+ : vector<128x64x64xf16>, vector<64x64xf16> into vector<128x64xf16>
+ return %3 : vector<128x64xf16>
+}
+
+// CHECK-LABEL: @negative_non_batched_nd(
+// CHECK: vector.contract
>From 6a2f80e02ff9dedf7220f0381050c05e25b045fc Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 25 Jul 2026 02:50:46 +0000
Subject: [PATCH 2/2] fix format
---
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 850fd80793736..29f5fb3a1d9b2 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -1001,9 +1001,10 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
if (!getRowMajorMatmulBatchRank(contractOp.getIndexingMapsAttr()))
return rewriter.notifyMatchFailure(
- contractOp, "Expects a (batched) row-major matmul: leading dims must "
- "be batch dims shared by lhs, rhs, and acc; innermost two "
- "dims must be (M, K), (K, N), and (M, N)");
+ contractOp,
+ "Expects a (batched) row-major matmul: leading dims must "
+ "be batch dims shared by lhs, rhs, and acc; innermost two "
+ "dims must be (M, K), (K, N), and (M, N)");
// xegpu.dpas operands are limited to rank 4 (2 batch + 2 core dims).
if (accType.getRank() > 4)
More information about the Mlir-commits
mailing list