shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir][tosa] Fix select folder when operands are broadcast (#165481)

Browse Source 
This commit addresses a crash in the dialects folder. The currently
folder assumes no broadcasting of the input operand happens and
therefore the folder can complain that the returned value was not the
same
shape as the result.

For now, this commit ensures no folding happens when broadcasting is
involved. In the future, folding with a broadcast could likely be
supported by inserting a `tosa.tile` operation before returning the
operand. This type of transformation is likely better suited for a
canonicalization pass. This commit only aims to avoid the crash.
This commit is contained in:
Luke Hutton 2025-11-20 14:39:28 +00:00 committed by GitHub
parent 364fe55c42
commit a74bfc06e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 59 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
View File

@ -1485,7 +1485,24 @@ OpFoldResult SliceOp::fold(FoldAdaptor adaptor) {
return {};
}
static bool
mayRequireBroadcast(ValueTypeRange<mlir::OperandRange> operandTypes) {
const auto isDynamic = [](Type ty) {
const auto shapedTy = llvm::dyn_cast<ShapedType>(ty);
return !shapedTy || !shapedTy.hasStaticShape();
};
return llvm::any_of(operandTypes, isDynamic) ||
failed(verifyCompatibleShapes(operandTypes));
}
OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
// Select allows operand shapes to be broadcast to the output shape. For
// now, don't support folding when we cannot prove no broadcasting is
// involved.
if (mayRequireBroadcast(getOperandTypes()))
return {};
if (getOnTrue() == getOnFalse())
return getOnTrue();

42
mlir/test/Dialect/Tosa/canonicalize.mlir
View File

@ -643,6 +643,48 @@ func.func @select_not_pred(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2:
// -----
// CHECK-LABEL: @select_broadcast_same_value_no_fold
func.func @select_broadcast_same_value_no_fold(%arg0: tensor<2x2xi1>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
// CHECK: tosa.select %arg0, %arg1, %arg1
%0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
// -----
// CHECK-LABEL: @select_broadcast_true_value_no_fold
func.func @select_broadcast_true_value_no_fold(%arg0: tensor<1x1xf32>, %arg1: tensor<2x2xf32>) -> tensor<?x?xf32> {
// CHECK: %[[CONST:.*]] = "tosa.const"
%0 = "tosa.const"() {values = dense<1> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
// CHECK: tosa.select %[[CONST]], %arg0, %arg1
%1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<2x2xf32>) -> tensor<?x?xf32>
return %1 : tensor<?x?xf32>
}
// -----
// CHECK-LABEL: @select_broadcast_false_value_no_fold
func.func @select_broadcast_false_value_no_fold(%arg0: tensor<2x2xf32>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
// CHECK: %[[CONST:.*]] = "tosa.const"
%0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
// CHECK: tosa.select %[[CONST]], %arg0, %arg1
%1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x2xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
return %1 : tensor<2x2xf32>
}
// -----
// CHECK-LABEL: @select_broadcast_false_value_dynamic_operand_no_fold
func.func @select_broadcast_false_value_dynamic_operand_no_fold(%arg0: tensor<2x?xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
// CHECK: %[[CONST:.*]] = "tosa.const"
%0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
// CHECK: tosa.select %[[CONST]], %arg0, %arg1
%1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x?xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
return %1 : tensor<2x2xf32>
}
// -----
// CHECK-LABEL: @reduce_all_fold
func.func @reduce_all_fold(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
// CHECK: return %arg0