[Mlir-commits] [mlir] ea6d5c6 - [MLIR][XeGPU] Fix blocking pass for scf.if distribution (#207060)

[llvmlistbot at llvm.org]

Author: Nishant Patel
Date: 2026-07-08T10:40:47+02:00
New Revision: ea6d5c6a2e37705688c74167990115e535058065

URL: https://github.com/llvm/llvm-project/commit/ea6d5c6a2e37705688c74167990115e535058065
DIFF: https://github.com/llvm/llvm-project/commit/ea6d5c6a2e37705688c74167990115e535058065.diff

LOG: [MLIR][XeGPU] Fix blocking pass for scf.if distribution (#207060)

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
    mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
    mlir/test/Dialect/XeGPU/xegpu-blocking.mlir

Removed: 
    


################################################################################
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..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
@@ -924,6 +925,19 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
         recordTypes(init, {arg, res, yieldVal});
       return;
     }
+    if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
+      // 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())) {
+        SmallVector<Value> dests{res, thenYield.getOperand(idx)};
+        if (elseYield)
+          dests.push_back(elseYield.getOperand(idx));
+        recordTypes(res, dests);
+      }
+      return;
+    }
   });
   return loopArgTypes;
 }

diff  --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index f2c25f45f45bc..b0b16c2adba6e 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -815,3 +815,31 @@ gpu.module @test_kernel {
     gpu.return
   }
 }
+
+// -----
+// 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]>
+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 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 dense<0.0> : vector<16x64xf32>
+      %b = arith.constant dense<0.0> : vector<16xf32>
+      scf.yield %a, %b : vector<16x64xf32>, vector<16xf32>
+    }
+    %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
+  }
+}