[Mlir-commits] [mlir] [MLIR][XeGPU] Fix blocking pass for scf.if distribution (PR #207060)
Nishant Patel
llvmlistbot at llvm.org
Tue Jul 7 17:22:20 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/207060
>From d81f090b2cb072e29317ebac50070010344492bd Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 30 Jun 2026 17:59:49 +0000
Subject: [PATCH 1/3] Fix blocking for scf.if with 1:N expanded results
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 2 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 17 +++++++++
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 35 +++++++++++++++++++
3 files changed, 53 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 4f81085c21e86..57b8ac73f41ac 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -430,7 +430,7 @@ void XeGPUBlockingPass::runOnOperation() {
// `layout_result_N` lands on the wrong (renumbered) result, corrupting the
// count invariant and leaving the loop illegal.
op->walk([](Operation *loopOp) {
- if (!isa<scf::ForOp, scf::WhileOp, scf::ConditionOp>(loopOp))
+ if (!isa<scf::ForOp, scf::WhileOp, scf::ConditionOp, scf::IfOp>(loopOp))
return;
SmallVector<StringRef> toRemove;
for (const NamedAttribute &attr : loopOp->getAttrs()) {
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 1aad2aa77741b..9ab0d55f51be4 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -924,6 +924,23 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
recordTypes(init, {arg, res, yieldVal});
return;
}
+ if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
+ // scf.if is not loop-carried, but its results still expand 1:N during
+ // blocking. Record each result together with the then/else yield operands
+ // that feed it so they all convert to the same distributed type by Value
+ // identity. The yield operand's layout is the authoritative source (the
+ // result's own per-position layout attr is stripped before conversion).
+ scf::YieldOp thenYield = ifOp.thenYield();
+ scf::YieldOp elseYield = ifOp.elseBlock() ? ifOp.elseYield() : nullptr;
+ for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) {
+ Value thenVal = thenYield.getOperand(idx);
+ SmallVector<Value> dests{res, thenVal};
+ if (elseYield)
+ dests.push_back(elseYield.getOperand(idx));
+ recordTypes(thenVal, dests);
+ }
+ return;
+ }
});
return loopArgTypes;
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index f2c25f45f45bc..b94119fa33121 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -815,3 +815,38 @@ gpu.module @test_kernel {
gpu.return
}
}
+
+// -----
+// Test that an scf.if with results that expand 1:N during blocking is
+// legalized. The two results use different layouts (a regular 2D layout that
+// expands 1:N and a slice layout), and the slice result is consumed by a
+// broadcast/transpose chain. This exercises the scf.if structural type
+// conversion, which previously failed to legalize because the per-position
+// layout attrs were not stripped and the results were not recorded in the
+// loop-arg type map.
+#if_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#if_b = #xegpu.slice<#xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>, dims = [0]>
+gpu.module @test_kernel {
+ // CHECK-LABEL: func @if_one_to_n_results
+ gpu.func @if_one_to_n_results(%c: i1, %p: i64) {
+ %c0 = arith.constant 0 : index
+ // CHECK: scf.if
+ // CHECK-COUNT-8: vector<8x16xf32>
+ %r:2 = scf.if %c -> (vector<16x64xf32>, vector<16xf32>) {
+ %a = arith.constant {layout_result_0 = #if_a} dense<1.0> : vector<16x64xf32>
+ %b = arith.constant {layout_result_0 = #if_b} dense<1.0> : vector<16xf32>
+ scf.yield %a, %b : vector<16x64xf32>, vector<16xf32>
+ } else {
+ %a = arith.constant {layout_result_0 = #if_a} dense<0.0> : vector<16x64xf32>
+ %b = arith.constant {layout_result_0 = #if_b} dense<0.0> : vector<16xf32>
+ scf.yield %a, %b : vector<16x64xf32>, vector<16xf32>
+ } {layout_result_0 = #if_a, layout_result_1 = #if_b}
+ %bc = vector.broadcast %r#1 {layout_result_0 = #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>} : vector<16xf32> to vector<64x16xf32>
+ %tp = vector.transpose %bc, [1, 0] {layout_result_0 = #if_a} : vector<64x16xf32> to vector<16x64xf32>
+ %add = arith.addf %r#0, %tp {layout_result_0 = #if_a} : vector<16x64xf32>
+ %tr = arith.truncf %add {layout_result_0 = #if_a} : vector<16x64xf32> to vector<16x64xf16>
+ %td = xegpu.create_nd_tdesc %p, shape:[16, 64], strides:[64, 1] : i64 -> !xegpu.tensor_desc<16x64xf16, #if_a>
+ xegpu.store_nd %tr, %td[%c0, %c0] <{layout = #if_a}> : vector<16x64xf16>, !xegpu.tensor_desc<16x64xf16, #if_a>
+ gpu.return
+ }
+}
>From c5ceb9208dd9f0f99b42dd8d2e6d009d43a3753f Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 1 Jul 2026 17:18:20 +0000
Subject: [PATCH 2/3] Fix comments
---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 8 +++---
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 27 ++++++++-------------
2 files changed, 13 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 9ab0d55f51be4..7faada2b8223f 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -925,11 +925,9 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
return;
}
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
- // scf.if is not loop-carried, but its results still expand 1:N during
- // blocking. Record each result together with the then/else yield operands
- // that feed it so they all convert to the same distributed type by Value
- // identity. The yield operand's layout is the authoritative source (the
- // result's own per-position layout attr is stripped before conversion).
+ // scf.if results expand 1:N during blocking. Record each result with the
+ // then/else yield operands that feed it so they convert to the same
+ // distributed type by Value identity, using the yield operand's layout.
scf::YieldOp thenYield = ifOp.thenYield();
scf::YieldOp elseYield = ifOp.elseBlock() ? ifOp.elseYield() : nullptr;
for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) {
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index b94119fa33121..b0b16c2adba6e 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -817,15 +817,8 @@ gpu.module @test_kernel {
}
// -----
-// Test that an scf.if with results that expand 1:N during blocking is
-// legalized. The two results use different layouts (a regular 2D layout that
-// expands 1:N and a slice layout), and the slice result is consumed by a
-// broadcast/transpose chain. This exercises the scf.if structural type
-// conversion, which previously failed to legalize because the per-position
-// layout attrs were not stripped and the results were not recorded in the
-// loop-arg type map.
+// Test that an scf.if whose results expand 1:N during blocking is legalized.
#if_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
-#if_b = #xegpu.slice<#xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>, dims = [0]>
gpu.module @test_kernel {
// CHECK-LABEL: func @if_one_to_n_results
gpu.func @if_one_to_n_results(%c: i1, %p: i64) {
@@ -833,18 +826,18 @@ gpu.module @test_kernel {
// CHECK: scf.if
// CHECK-COUNT-8: vector<8x16xf32>
%r:2 = scf.if %c -> (vector<16x64xf32>, vector<16xf32>) {
- %a = arith.constant {layout_result_0 = #if_a} dense<1.0> : vector<16x64xf32>
- %b = arith.constant {layout_result_0 = #if_b} dense<1.0> : vector<16xf32>
+ %a = arith.constant dense<1.0> : vector<16x64xf32>
+ %b = arith.constant dense<1.0> : vector<16xf32>
scf.yield %a, %b : vector<16x64xf32>, vector<16xf32>
} else {
- %a = arith.constant {layout_result_0 = #if_a} dense<0.0> : vector<16x64xf32>
- %b = arith.constant {layout_result_0 = #if_b} dense<0.0> : vector<16xf32>
+ %a = arith.constant dense<0.0> : vector<16x64xf32>
+ %b = arith.constant dense<0.0> : vector<16xf32>
scf.yield %a, %b : vector<16x64xf32>, vector<16xf32>
- } {layout_result_0 = #if_a, layout_result_1 = #if_b}
- %bc = vector.broadcast %r#1 {layout_result_0 = #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>} : vector<16xf32> to vector<64x16xf32>
- %tp = vector.transpose %bc, [1, 0] {layout_result_0 = #if_a} : vector<64x16xf32> to vector<16x64xf32>
- %add = arith.addf %r#0, %tp {layout_result_0 = #if_a} : vector<16x64xf32>
- %tr = arith.truncf %add {layout_result_0 = #if_a} : vector<16x64xf32> to vector<16x64xf16>
+ }
+ %bc = vector.broadcast %r#1 : vector<16xf32> to vector<64x16xf32>
+ %tp = vector.transpose %bc, [1, 0] : vector<64x16xf32> to vector<16x64xf32>
+ %add = arith.addf %r#0, %tp : vector<16x64xf32>
+ %tr = arith.truncf %add : vector<16x64xf32> to vector<16x64xf16>
%td = xegpu.create_nd_tdesc %p, shape:[16, 64], strides:[64, 1] : i64 -> !xegpu.tensor_desc<16x64xf16, #if_a>
xegpu.store_nd %tr, %td[%c0, %c0] <{layout = #if_a}> : vector<16x64xf16>, !xegpu.tensor_desc<16x64xf16, #if_a>
gpu.return
>From 5e3ca5318a8e9e62f3c3a6de2214102b54ac06a8 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 6 Jul 2026 21:34:40 +0000
Subject: [PATCH 3/3] Address feedback
---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 19 +++++++++----------
1 file changed, 9 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 7faada2b8223f..9620e21f9bfdf 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -862,12 +862,13 @@ bool xegpu::matchSplitDimExpansion(
//===----------------------------------------------------------------------===//
// Pre-computes distributed VectorType mappings for every value carried through
-// an SCF loop (scf.while, scf.for): block args (iter_args /
-// before-/after-args), loop results, and the terminator operands feeding them.
+// an SCF region-branch op (scf.while, scf.for, scf.if): block args (iter_args /
+// before-/after-args), op results, and the terminator operands feeding them.
// These positions share one logical value and must convert identically, so each
// is derived from a single source -- the layout of the feeding value (loop
-// init, or `scf.condition` operand) -- via `getDistributeLayoutAttr(Value)`,
-// and keyed by `Value`. Keying by Value is required because the SCF converters
+// init, `scf.condition` operand, or `scf.if` result) -- via
+// `getDistributeLayoutAttr(Value)`, and keyed by `Value`. Keying by Value is
+// required because the SCF converters
// detach/replace the loop body mid-conversion (scf.while detaches before/after
// blocks -> a detached-arg layout query trips an ilist assertion; scf.for
// rebuilds the op, which loses the temporary `layout_operand_N` attrs -> the
@@ -925,17 +926,15 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
return;
}
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
- // scf.if results expand 1:N during blocking. Record each result with the
- // then/else yield operands that feed it so they convert to the same
- // distributed type by Value identity, using the yield operand's layout.
+ // Each result and its then/else yield operands share one position and
+ // must convert identically; derive all from the result's layout.
scf::YieldOp thenYield = ifOp.thenYield();
scf::YieldOp elseYield = ifOp.elseBlock() ? ifOp.elseYield() : nullptr;
for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) {
- Value thenVal = thenYield.getOperand(idx);
- SmallVector<Value> dests{res, thenVal};
+ SmallVector<Value> dests{res, thenYield.getOperand(idx)};
if (elseYield)
dests.push_back(elseYield.getOperand(idx));
- recordTypes(thenVal, dests);
+ recordTypes(res, dests);
}
return;
}
More information about the Mlir-commits
mailing list