Yield if condition, range-based escaping values for innerwarps

Author: akroviakov

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

@@ -1794,14 +1794,18 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
       return failure();
 
     // The new `WarpOp` groups yields values in following order:
-    // 1. Escaping values then branch
-    // 2. Escaping values else branch
-    // 3. All non-`ifOp` yielded values.
-    SmallVector<Value> newWarpOpYieldValues{escapingValuesThen.begin(),
-                                            escapingValuesThen.end()};
+    // 1. Branch condition
+    // 2. Escaping values then branch
+    // 3. Escaping values else branch
+    // 4. All non-`ifOp` yielded values.
+    SmallVector<Value> newWarpOpYieldValues{ifOp.getCondition()};
+    newWarpOpYieldValues.append(escapingValuesThen.begin(),
+                                escapingValuesThen.end());
     newWarpOpYieldValues.append(escapingValuesElse.begin(),
                                 escapingValuesElse.end());
-    SmallVector<Type> newWarpOpDistTypes = escapingValueDistTypesThen;
+    SmallVector<Type> newWarpOpDistTypes{ifOp.getCondition().getType()};
+    newWarpOpDistTypes.append(escapingValueDistTypesThen.begin(),
+                              escapingValueDistTypesThen.end());
     newWarpOpDistTypes.append(escapingValueDistTypesElse.begin(),
                               escapingValueDistTypesElse.end());
 
@@ -1815,7 +1819,6 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
     // Create the new `WarpOp` with the updated yield values and types.
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndReplaceReturns(
         rewriter, warpOp, newWarpOpYieldValues, newWarpOpDistTypes);
-
     // `ifOp` returns the result of the inner warp op.
     SmallVector<Type> newIfOpDistResTypes;
     for (auto [i, res] : llvm::enumerate(ifOp.getResults())) {
@@ -1831,23 +1834,24 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
     // Create a new `IfOp` outside the new `WarpOp` region.
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPointAfter(newWarpOp);
-    auto newIfOp = scf::IfOp::create(rewriter, ifOp.getLoc(),
-                                     newIfOpDistResTypes, ifOp.getCondition(),
-                                     static_cast<bool>(ifOp.thenBlock()),
-                                     static_cast<bool>(ifOp.elseBlock()));
+    auto newIfOp = scf::IfOp::create(
+        rewriter, ifOp.getLoc(), newIfOpDistResTypes, newWarpOp.getResult(0),
+        static_cast<bool>(ifOp.thenBlock()),
+        static_cast<bool>(ifOp.elseBlock()));
 
     auto processBranch = [&](Block *oldIfBranch, Block *newIfBranch,
                              llvm::SmallSetVector<Value, 32> &escapingValues,
-                             SmallVector<Type> &escapingValueInputTypes) {
+                             SmallVector<Type> &escapingValueInputTypes,
+                             size_t warpResRangeStart) {
       OpBuilder::InsertionGuard g(rewriter);
       if (!newIfBranch)
         return;
       rewriter.setInsertionPointToStart(newIfBranch);
       llvm::SmallDenseMap<Value, int64_t> escapeValToBlockArgIndex;
       SmallVector<Value> innerWarpInputVals;
       SmallVector<Type> innerWarpInputTypes;
-      for (size_t i = 0; i < escapingValues.size(); ++i) {
-        innerWarpInputVals.push_back(newWarpOp.getResult(i));
+      for (size_t i = 0; i < escapingValues.size(); ++i, ++warpResRangeStart) {
+        innerWarpInputVals.push_back(newWarpOp.getResult(warpResRangeStart));
         escapeValToBlockArgIndex[escapingValues[i]] =
             innerWarpInputTypes.size();
         innerWarpInputTypes.push_back(escapingValueInputTypes[i]);
@@ -1886,11 +1890,11 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
     };
     processBranch(&ifOp.getThenRegion().front(),
                   &newIfOp.getThenRegion().front(), escapingValuesThen,
-                  escapingValueInputTypesThen);
+                  escapingValueInputTypesThen, 1);
     if (!ifOp.getElseRegion().empty())
       processBranch(&ifOp.getElseRegion().front(),
                     &newIfOp.getElseRegion().front(), escapingValuesElse,
-                    escapingValueInputTypesElse);
+                    escapingValueInputTypesElse, 1 + escapingValuesThen.size());
     // Update the users of `<- WarpOp.yield <- IfOp.yield` to use the new `IfOp`
     // result.
     for (auto [origIdx, newIdx] : ifResultMapping)

mlir/test/Dialect/XeGPU/subgroup-distribute.mlir

@@ -339,11 +339,11 @@ gpu.module @test {
 }
 
 // -----
-// CHECK-LABEL: gpu.func @scatter_ops_scf_yield({{.*}}) {
+// CHECK-LABEL: gpu.func @scatter_ops_scf_yield({{.*}},
+// CHECK-SAME: %[[PREDICATE:.*]]: i1) {
 // CHECK: %[[DEFAULT:.*]] = arith.constant dense<1.200000e+01> : vector<8xf16>
 // CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
 // CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[PREDICATE:.*]] = llvm.mlir.poison : i1
 // CHECK: %[[PREDICATED_LOAD:.*]] = scf.if %[[PREDICATE]] -> (vector<8xf16>) {
 // CHECK-NEXT: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}> : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
 // CHECK-NEXT: scf.yield %[[LOADED]] : vector<8xf16>
@@ -352,8 +352,7 @@ gpu.module @test {
 // CHECK-NEXT: }
 // CHECK-NEXT: xegpu.store %[[PREDICATED_LOAD]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
 gpu.module @test {
-  gpu.func @scatter_ops_scf_yield(%src: memref<256xf16>) {
-    %pred = llvm.mlir.poison : i1
+  gpu.func @scatter_ops_scf_yield(%src: memref<256xf16>, %pred : i1) {
     %1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1>: vector<16xi1>
     %offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
     %loaded = scf.if %pred -> (vector<16x8xf16>) {