[Mlir-commits] [mlir] [mlir][linalg] Fix getSourceSkipUnary to only skip cast-like ops (PR #198725)

[Hocky Yudhiono]

================
@@ -46,3 +46,48 @@ func.func @not_copy(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32
   } -> tensor<8xi32>
   return %res : tensor<8xi32>
 }
+
+// -----
+
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// This test checks that linalg.generic with a negf between mulf and addf
+// does not get incorrectly specialized to matmul.
+// CHECK-LABEL: @contraction_with_negf
+//  CHECK-NOT:    linalg.matmul
+//      CHECK:    linalg.generic
+func.func @contraction_with_negf(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
+  %0 = linalg.generic {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %1 = arith.mulf %in, %in_0 : f32
+    %2 = arith.negf %1 : f32
+    %3 = arith.addf %out, %2 : f32
+    linalg.yield %3 : f32
+  } -> tensor<3x3xf32>
+  return %0 : tensor<3x3xf32>
+}
+
+// -----
+
+#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// This test checks that a cast chain changing the input semantics does not get
+// ignored when matching contractions.
+// CHECK-LABEL: @contraction_with_rounding_cast_chain
+//  CHECK-NOT:    linalg.matmul
+//      CHECK:    linalg.generic
+func.func @contraction_with_rounding_cast_chain(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
+  %0 = linalg.generic {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %1 = arith.fptosi %in : f32 to i32
+    %2 = arith.sitofp %1 : i32 to f32
+    %3 = arith.fptosi %in_0 : f32 to i32
+    %4 = arith.sitofp %3 : i32 to f32
+    %5 = arith.mulf %2, %4 : f32
+    %6 = arith.addf %out, %5 : f32
+    linalg.yield %6 : f32
+  } -> tensor<3x3xf32>
+  return %0 : tensor<3x3xf32>
+}
----------------
hockyy wrote:

can we add a failing testcase for linalg.contract as well

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