[Mlir-commits] [mlir] [mlir][tosa] Improve broadcasting behaviour in elementwise folders (PR #181114)

[Tai Ly]

================
@@ -136,6 +136,56 @@ func.func @fold_add_splat_ui8_overflow() -> tensor<10xui8> {
 
 // -----
 
+// CHECK-LABEL: @no_fold_add_unknown_broadcast_zero_lhs
+func.func @no_fold_add_unknown_broadcast_zero_lhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
+  %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+  %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<1x4xi32>) -> tensor<2x4xi32>
+  // CHECK: tosa.add
+  return %div : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_lhs
+func.func @no_fold_dynamic_add_unknown_broadcast_zero_lhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+  %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+  %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
+  // CHECK: tosa.add
+  return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_lhs
+func.func @fold_dynamic_add_broadcast_zero_lhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+  %one = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+  %div = tosa.add %one, %arg0 : (tensor<1x1xi32>, tensor<?x17xi32>) -> tensor<?x17xi32>
+  // CHECK: return %arg0
+  return %div : tensor<?x17xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_rhs
+func.func @no_fold_dynamic_add_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+  %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+  %div = tosa.add %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+  // CHECK: tosa.add
+  return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_rhs
+func.func @fold_dynamic_add_broadcast_zero_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+  %one = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
----------------
Tai78641 wrote:

ditto

https://github.com/llvm/llvm-project/pull/181114