[Mlir-commits] [mlir] [mlir][linalg] Avoid asserts in IndexingMapOpInterface (PR #179072)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

Replace the failsafe assertions with conservative fallbacks: return 0 rank / empty shape when the operand is not a ranked tensor or memref. 

Fixes https://github.com/llvm/llvm-project/issues/179043


---
Full diff: https://github.com/llvm/llvm-project/pull/179072.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td (+6-6) 
- (modified) mlir/test/Dialect/Linalg/invalid.mlir (+19) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 9f1e88a040f5f..365d5344953b6 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -385,9 +385,9 @@ def LinalgStructuredInterface
           return 0;
         // Tensor and Memref container types have a rank.
         if (auto shapedType = ::llvm::dyn_cast<ShapedType>(t)) {
-          // Failsafe.
-          assert((isa<MemRefType>(t) || isa<RankedTensorType>(t)) &&
-                 "expected a ranked tensor or memref in LinalgInterface::getRank");
+          // Only ranked tensors and MemRefs have well defined ranks.
+          if (!(isa<MemRefType>(t) || isa<RankedTensorType>(t)))
+            return 0;
           return shapedType.getRank();
         }
         return 0;
@@ -703,9 +703,9 @@ def LinalgStructuredInterface
         if (isa<VectorType>(t))
           return {};
         if (auto shapedType = ::llvm::dyn_cast<ShapedType>(t)) {
-          // Failsafe.
-          assert((isa<MemRefType>(t) || isa<RankedTensorType>(t)) &&
-                 "expected a ranked tensor or memref in LinalgInterface::getRank");
+          // Only ranked tensors and MemRefs have well defined shapes.
+          if (!(isa<MemRefType>(t) || isa<RankedTensorType>(t)))
+            return {};
           return shapedType.getShape();
         }
         return {};
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 355d801f8732c..d97b5ee31f97b 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -217,6 +217,25 @@ func.func @generic_indexing_map_with_symbol(%arg0: tensor<8xf32>) -> tensor<8xf3
   return %0 : tensor<8xf32>
 }
 
+// -----
+
+// Unranked tensor inputs must be diagnosed.
+func.func @generic_unranked_input_tensor(%in: tensor<*xf32>) {
+  %out = tensor.empty() : tensor<16x16xf32>
+  // expected-error @+1 {{expected operand #0 rank (0) to match the result rank of indexing_map (2)}}
+  %r = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                     affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]}
+      ins(%in : tensor<*xf32>)
+     outs(%out : tensor<16x16xf32>) {
+      ^bb0(%a: f32, %b: f32):
+        linalg.yield %a : f32
+     } -> tensor<16x16xf32>
+  return
+}
+
+
 ////////////////////////////////////////////////////////////////////////////////
 ///////////////////////////// Region tests /////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////

``````````

</details>


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