[Mlir-commits] [mlir] [MLIR][XeGPU] Fix blocking pass for scf.if distribution (PR #207060)
Nishant Patel
llvmlistbot at llvm.org
Wed Jul 1 12:06:36 PDT 2026
https://github.com/nbpatel created https://github.com/llvm/llvm-project/pull/207060
None
>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/2] 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/2] 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
More information about the Mlir-commits
mailing list