Skip to content

[mlir][vector] Fix for WarpOpScfForOp failure when scf.for has results that are unused. #141853

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Jun 9, 2025
+65 −10
Merged

[mlir][vector] Fix for WarpOpScfForOp failure when scf.for has results that are unused. #141853

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
ff1012e
add bug fix
charithaintc May 27, 2025
c6eb53f
add test
charithaintc May 28, 2025
3bdb596
add comments
charithaintc May 28, 2025
6d47e3f
Merge branch 'main' into scf_for_bug
charithaintc May 28, 2025
fe3ab99
remove unsused headers
charithaintc May 28, 2025
cc621a1
Merge branch 'main' into scf_for_bug
charithaintc Jun 2, 2025
e1a1bc1
Merge branch 'main' into scf_for_bug
charithaintc Jun 4, 2025
7d25981
fix
charithaintc Jun 4, 2025
e883ba4
Merge branch 'main' into scf_for_bug
charithaintc Jun 5, 2025
3ae7706
Merge branch 'main' into scf_for_bug
charithaintc Jun 9, 2025
5224464
Merge branch 'main' into scf_for_bug
charithaintc Jun 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 30 additions & 10 deletions mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1554,22 +1554,37 @@ struct WarpOpScfForOp : public WarpDistributionPattern {
llvm::SmallSetVector<Value, 32> escapingValues;
SmallVector<Type> inputTypes;
SmallVector<Type> distTypes;
auto collectEscapingValues = [&](Value value) {
if (!escapingValues.insert(value))
return;
Type distType = value.getType();
if (auto vecType = dyn_cast<VectorType>(distType)) {
AffineMap map = distributionMapFn(value);
distType = getDistributedType(vecType, map, warpOp.getWarpSize());
}
inputTypes.push_back(value.getType());
distTypes.push_back(distType);
};

mlir::visitUsedValuesDefinedAbove(
forOp.getBodyRegion(), [&](OpOperand *operand) {
Operation *parent = operand->get().getParentRegion()->getParentOp();
if (warpOp->isAncestor(parent)) {
if (!escapingValues.insert(operand->get()))
return;
Type distType = operand->get().getType();
if (auto vecType = dyn_cast<VectorType>(distType)) {
AffineMap map = distributionMapFn(operand->get());
distType = getDistributedType(vecType, map, warpOp.getWarpSize());
}
inputTypes.push_back(operand->get().getType());
distTypes.push_back(distType);
collectEscapingValues(operand->get());
}
});

// Any forOp result that is not already yielded by the warpOp
// region is also considered escaping and must be returned by the
// original warpOp.
for (OpResult forResult : forOp.getResults()) {
// Check if this forResult is already yielded by the yield op.
if (llvm::is_contained(yield->getOperands(), forResult)) {
continue;
}
collectEscapingValues(forResult);
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't get this part, could you please clarify? This should collect the arguments of the scf.for that would be distributed and returned by the new signature of the original warp op. Why would the original warp op return scf.for's results?
Shouldn't the results be only collected and inserted into the second warp op? The one that is created inside scf.for.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great question. This is how the original implementation work. Initially original warpOp will return the used forOp results + any escaping values. Later the for op results are replaced with the corresponding initArgs, Here:

llvm-project/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

Line 1607 in fe3ab99

yieldOperand.set(forOp.getInitArgs()[forResult.getResultNumber()]);

I did not change the original logic flow. Hope this clarifies your question.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, I had to play with it for a bit to understand how it works. I think adding results of the ForOp to the yielded values as the first step is pretty confusing since the resulting warp op after transformation should not have those. I've modified your lit to make an example that would probably be helpful to others to understand it better.
Initial IR:

func.func @warp_scf_for_unused_yield(%arg0: index) {
  %c128 = arith.constant 128 : index
  %c1 = arith.constant 1 : index
  %c0 = arith.constant 0 : index
  %0:2 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>) {
    %ini = "some_def"() : () -> (vector<128xf32>)
    %ini1 = "some_def"() : () -> (vector<128xf32>)
    %other = "other_def"() : () -> (vector<128xf32>)
    %3:3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini, %arg5 = %ini1, %other_2 = %other) -> (vector<128xf32>, vector<128xf32>, vector<128xf32>) {
      %add = arith.addi %arg3, %c1 : index
      %1  = "some_def"(%arg5, %add) : (vector<128xf32>, index) -> (vector<128xf32>)
      %acc = "some_def"(%add, %arg4, %1) : (index, vector<128xf32>, vector<128xf32>) -> (vector<128xf32>)
      %other2 = "other2_def"(%arg4) : (vector<128xf32>) -> (vector<128xf32>)
      scf.yield %acc, %1, %other2 : vector<128xf32>, vector<128xf32>, vector<128xf32>
    }
    gpu.yield %3#0, %other : vector<128xf32>, vector<128xf32>
  }
  "some_use"(%0#0) : (vector<4xf32>) -> ()
  "other_use"(%0#1) : (vector<4xf32>) -> ()
  return
}

Right after the newWarpOp creation the warp op looks like:

%0:4 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) {
  %1 = "some_def"() : () -> vector<128xf32>
  %2 = "some_def"() : () -> vector<128xf32>
  %3 = "other_def"() : () -> vector<128xf32>
  %4:3 = scf.for %arg1 = %c0 to %c128 step %c1 iter_args(%arg2 = %1, %arg3 = %2, %arg4 = %3) -> (vector<128xf32>, vector<128xf32>, vector<128xf32>) {
    %5 = arith.addi %arg1, %c1 : index
    %6 = "some_def"(%arg3, %5) : (vector<128xf32>, index) -> vector<128xf32>
    %7 = "some_def"(%5, %arg2, %6) : (index, vector<128xf32>, vector<128xf32>) -> vector<128xf32>
    %8 = "other2_def"(%arg2) : (vector<128xf32>) -> vector<128xf32>
    scf.yield %7, %6, %8 : vector<128xf32>, vector<128xf32>, vector<128xf32>
  }
  gpu.yield %4#0, %3, %4#1, %4#2 : vector<128xf32>, vector<128xf32>, vector<128xf32>, vector<128xf32>
}

The results of the ForOp are now returned even though the op itself will be sinked through gpu.yield. These are patched up later again to produce the valid result. This is a bit unexpected after you read the commend for the transformation (and the reason it caught my attention).

I think this is somewhat brittle. If it is done to preserve some existing code I'd reconsider in favor of simplicity.

Copy link
Contributor Author

@charithaintc charithaintc May 30, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review :-)

I agree the existing logic is bit confusing. I am also not sure why the forOp results are yielded instead of the iterArgs directly. Maybe code owners have some insights on this?

But overall the logic will be similar apart from this part (even for a refactored version). This pattern is anyway bit complex since it has a lot of moving parts. :-)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the code is to add the initial value of the loop-carried variables to the new warp op's result -

Instead collect the loop result and fixup latter, can we directly collect "forOp.getInitArgs()" as escaping value - since these are escaping value of the for statement, but not captured by the original code which only scan the loop body.

if (llvm::is_contained(distTypes, Type{}))
return failure();

Expand Down Expand Up @@ -1609,7 +1624,12 @@ struct WarpOpScfForOp : public WarpDistributionPattern {
forOp.getResultTypes().end());
llvm::SmallDenseMap<Value, int64_t> argIndexMapping;
for (auto [i, retIdx] : llvm::enumerate(newRetIndices)) {
warpInput.push_back(newWarpOp.getResult(retIdx));
auto newWarpResult = newWarpOp.getResult(retIdx);
// Unused forOp results yielded by the warpOp region are already included
// in the new ForOp.
if (llvm::is_contained(newOperands, newWarpResult))
continue;
warpInput.push_back(newWarpResult);
argIndexMapping[escapingValues[i]] = warpInputType.size();
warpInputType.push_back(inputTypes[i]);
}
Expand Down
36 changes: 36 additions & 0 deletions mlir/test/Dialect/Vector/vector-warp-distribute.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -584,6 +584,42 @@ func.func @warp_scf_for_multiple_yield(%arg0: index, %arg1: memref<?xf32>, %arg2
return
}

// -----
// CHECK-PROP-LABEL: func.func @warp_scf_for_unused_yield(
// CHECK-PROP: %[[W0:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4xf32>) {
// CHECK-PROP: %[[INI0:.*]] = "some_def"() : () -> vector<128xf32>
// CHECK-PROP: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>
// CHECK-PROP: gpu.yield %[[INI0]], %[[INI1]] : vector<128xf32>, vector<128xf32>
// CHECK-PROP: }
// CHECK-PROP: %[[F:.*]]:2 = scf.for %{{.*}} iter_args(%{{.*}} = %[[W0]]#0, %{{.*}} = %[[W0]]#1) -> (vector<4xf32>, vector<4xf32>) {
// CHECK-PROP: %[[W1:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%{{.*}} : vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) {
// CHECK-PROP: %[[ACC0:.*]] = "some_def"(%{{.*}}) : (vector<128xf32>, index) -> vector<128xf32>
// CHECK-PROP: %[[ACC1:.*]] = "some_def"(%{{.*}}) : (index, vector<128xf32>, vector<128xf32>) -> vector<128xf32>
// CHECK-PROP: gpu.yield %[[ACC1]], %[[ACC0]] : vector<128xf32>, vector<128xf32>
// CHECK-PROP: }
// CHECK-PROP: scf.yield %[[W1]]#0, %[[W1]]#1 : vector<4xf32>, vector<4xf32>
// CHECK-PROP: }
// CHECK-PROP: "some_use"(%[[F]]#0) : (vector<4xf32>) -> ()
func.func @warp_scf_for_unused_yield(%arg0: index) {
%c128 = arith.constant 128 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) {
%ini = "some_def"() : () -> (vector<128xf32>)
%ini1 = "some_def"() : () -> (vector<128xf32>)
%3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini, %arg5 = %ini1) -> (vector<128xf32>, vector<128xf32>) {
%add = arith.addi %arg3, %c1 : index
%1 = "some_def"(%arg5, %add) : (vector<128xf32>, index) -> (vector<128xf32>)
%acc = "some_def"(%add, %arg4, %1) : (index, vector<128xf32>, vector<128xf32>) -> (vector<128xf32>)
scf.yield %acc, %1 : vector<128xf32>, vector<128xf32>
}
gpu.yield %3#0 : vector<128xf32>
}
"some_use"(%0) : (vector<4xf32>) -> ()
return
}


// -----

// CHECK-PROP-LABEL: func @vector_reduction(
Expand Down
Loading