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

[Ayokunle Amodu]

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

>From 80200160d4c4288f604046440e6dd36a924f0fdc Mon Sep 17 00:00:00 2001
From: Ayokunle Amodu <121697771+ayokunle321 at users.noreply.github.com>
Date: Wed, 15 Jan 2025 15:32:35 -0700
Subject: [PATCH 1/4] added check for reduction op in contraction body

---
 mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index caf9cdb3a3eb4f..14f8f9e8fdd3b4 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -281,6 +281,12 @@ bool mlir::linalg::detail::isContractionBody(
 
   Value yielded = getSourceSkipUnary(terminator->getOperand(0));
   Operation *reductionOp = yielded.getDefiningOp();
+
+  if (!reductionOp){
+    errs << "expected reduction op in body";
+    return false;
+  }
+
   if (reductionOp->getNumResults() != 1 || reductionOp->getNumOperands() != 2) {
     errs << "expected reduction op to be binary";
     return false;

>From 4ed59274eebea0d0601b37d5b138d4ed57371617 Mon Sep 17 00:00:00 2001
From: Ayokunle Amodu <121697771+ayokunle321 at users.noreply.github.com>
Date: Wed, 15 Jan 2025 16:01:19 -0700
Subject: [PATCH 2/4] fix code style

---
 mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index 14f8f9e8fdd3b4..91165ddeb88870 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -282,7 +282,7 @@ bool mlir::linalg::detail::isContractionBody(
   Value yielded = getSourceSkipUnary(terminator->getOperand(0));
   Operation *reductionOp = yielded.getDefiningOp();
 
-  if (!reductionOp){
+  if (!reductionOp) {
     errs << "expected reduction op in body";
     return false;
   }

>From d72f50420e5118c1f26dbfff9c4eac4a2e872f0a Mon Sep 17 00:00:00 2001
From: Ayokunle Amodu <121697771+ayokunle321 at users.noreply.github.com>
Date: Thu, 16 Jan 2025 07:39:53 -0700
Subject: [PATCH 3/4] add test case to specialize-generic-ops-fail.mlir

---
 .../Linalg/specialize-generic-ops-fail.mlir   | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
index 542a7ed4a198b8..9993046f1325f1 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
@@ -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)>
+// Verifies that the pass crashes when trying to specialize a linalg.generic op with 
+// a reduction iterator when there is no valid reduction operation in the contraction 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>
+}
\ No newline at end of file

>From 9472c46c18ae32a905c6f53c10cf2d2e73bfd996 Mon Sep 17 00:00:00 2001
From: Ayokunle Amodu <121697771+ayokunle321 at users.noreply.github.com>
Date: Thu, 16 Jan 2025 07:51:38 -0700
Subject: [PATCH 4/4] edit test description, need to minimize test

---
 mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
index 9993046f1325f1..abe7a33c3a9e86 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
@@ -21,8 +21,8 @@ func.func @transpose_and_broadcast(%arg0: tensor<7x8xf32>, %arg1: tensor<8x7x9xf
 #map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
 #map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
 #map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// Verifies that the pass crashes when trying to specialize a linalg.generic op with 
-// a reduction iterator when there is no valid reduction operation in the contraction body.
+// 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