[Mlir-commits] [mlir] [mlir][linalg] Add Check for Reduction Operation in Contraction Body (PR #123134)

[Adam Siemieniuk]

================
@@ -14,3 +14,39 @@ func.func @transpose_and_broadcast(%arg0: tensor<7x8xf32>, %arg1: tensor<8x7x9xf
   } -> tensor<8x7x9xf32>
   return %0 : tensor<8x7x9xf32>
 }
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// This tests checks that the pass does not crash when trying to specialize a 
+// contraction-like generic op with no reduction operation in its body.
+// CHECK-LABEL: @specialize_reduction
+func.func private @specialize_reduction(%arg0: tensor<1x31x8xi32>) -> tensor<31x31xi32> {
+  %c-2351_i32 = arith.constant -2351 : i32
+  %c0_i32 = arith.constant 0 : i32
+  %0 = tensor.empty() : tensor<31x8xi32>
+  %1 = linalg.generic 
+        {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]}
+        outs(%0 : tensor<31x8xi32>) {
+    ^bb0(%out: i32):
+      linalg.yield %c-2351_i32 : i32
+  } -> tensor<31x8xi32>
+  %2 = tensor.empty() : tensor<31x31xi32>
+  %3 = linalg.generic 
+        {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]}
+        outs(%2 : tensor<31x31xi32>) {
+    ^bb0(%out: i32):
+      linalg.yield %c0_i32 : i32
+  } -> tensor<31x31xi32>
+  %4 = linalg.generic 
+        {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"]}
+        ins(%1, %1 : tensor<31x8xi32>, tensor<31x8xi32>) 
+        outs(%3 : tensor<31x31xi32>) {
+    ^bb0(%in: i32, %in_0: i32, %out: i32):
+      linalg.yield %out : i32
+  } -> tensor<31x31xi32>
+  return %4 : tensor<31x31xi32>
+}
----------------
adam-smnk wrote:

nit: missing newline at the end

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