[Mlir-commits] [mlir] [mlir][tosa] Fix select folder when operands are broadcast (PR #165481)

[Luke Hutton]

https://github.com/lhutton1 created https://github.com/llvm/llvm-project/pull/165481

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.

>From 2983504f8502011b714e3841bcd03519aeaf1790 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 27 Oct 2025 11:30:38 +0000
Subject: [PATCH] [mlir][tosa] Fix select folder when operands are broadcast

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.

Change-Id: I2a3a21a3d2a626d4dd013219de7c37d656977979
---
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 17 ++++++++
 mlir/test/Dialect/Tosa/canonicalize.mlir      | 42 +++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index a85ff10aa0d73..a4217051f2dfa 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1484,7 +1484,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();
 
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 5a40f3fa8572c..fc5ea7710e2c4 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -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