[mlir][linalg] Fix bug in InferStaticShapeOfOperands pattern

[vinograd47]

Summary

The pattern tries to deduce static shape from tensor.cast producer of linalg operation operands.
The original code unconditionally casts type of the tensor.cast source to RankedTensorType.
But the tensor.cast can also operate on UnrankedTensorType, so this cast either fail on assertion
in debug build or introduce UB in release build.

The patch replaces unconditional cast with dyn_cast and check for the cast result.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

for (OpOperand &opOperand : operands) {
     // Get the `sourceShape` of the `sourceType`. If the operand is a result of
     // `tensor.cast` operation and source of the cast operation has a static
     // shape, then assign it to the `sourceShape`.
     auto *parentOp = src.getDefiningOp();
     ArrayRef<int64_t> sourceShape = sourceType.getShape();
     if (parentOp) {
       if (auto castOp = dyn_cast<tensor::CastOp>(parentOp)) {
         Value castSource = castOp.getSource();
-        auto castSourceType = castSource.getType().cast<RankedTensorType>();
-        if (castSourceType.hasStaticShape())
+        auto castSourceType = castSource.getType().dyn_cast<RankedTensorType>();
+        if (castSourceType && castSourceType.hasStaticShape())
           sourceShape = castSourceType.getShape();
       }
     }
 
     // If the source shape's dimension has a static shape, map the affine dim
     // expression to the known static size.
     for (unsigned i = 0; i < sourceShape.size(); i++) {
       if (sourceType.isDynamicDim(i))

mlir/test/Dialect/Linalg/canonicalize.mlir

 // CHECK-LABEL: @dce_zero_memref
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<0xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<0xf32>
 //   CHECK-NOT:   memref.copy
 //  CHECK-NEXT:   return %[[ARG1]]
 
 // -----
 
-
 // CHECK-LABEL: func @tensor.cast(
 func.func @tensor.cast(%a : tensor<3x4xf32>, %b : tensor<4x?xf32>, %c : tensor<3x?xf32>)
   -> tensor<3x?xf32>
 {
   %ta = tensor.cast %a : tensor<3x4xf32> to tensor<?x?xf32>
   %tb = tensor.cast %b : tensor<4x?xf32> to tensor<?x?xf32>
   %tc = tensor.cast %c : tensor<3x?xf32> to tensor<?x?xf32>
 
   //      CHECK:  linalg.matmul ins({{.*}}tensor<3x4xf32>, tensor<4x?xf32>)
   // CHECK-SAME:    outs({{.*}}tensor<3x?xf32>) -> tensor<3x?xf32>
   %0 = linalg.matmul ins(%ta, %tb: tensor<?x?xf32>, tensor<?x?xf32>)
                     outs(%tc: tensor<?x?xf32>) -> tensor<?x?xf32>
 
   %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<3x?xf32>
 
   return %1: tensor<3x?xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @tensor.cast.unranked(
+func.func @tensor.cast.unranked(%a : tensor<*xf32>, %b : tensor<*xf32>, %c : tensor<*xf32>)
+  -> tensor<*xf32>
+{
+  //      CHECK:  tensor.cast
+  //      CHECK:  tensor.cast
+  //      CHECK:  tensor.cast
+  %ta = tensor.cast %a : tensor<*xf32> to tensor<?x?xf32>
+  %tb = tensor.cast %b : tensor<*xf32> to tensor<?x?xf32>
+  %tc = tensor.cast %c : tensor<*xf32> to tensor<?x?xf32>
+
+  //      CHECK:  linalg.matmul ins({{.*}}tensor<?x?xf32>, tensor<?x?xf32>)
+  // CHECK-SAME:    outs({{.*}}tensor<?x?xf32>) -> tensor<?x?xf32>
+  %0 = linalg.matmul ins(%ta, %tb: tensor<?x?xf32>, tensor<?x?xf32>)
+                    outs(%tc: tensor<?x?xf32>) -> tensor<?x?xf32>
+
+  //      CHECK:  tensor.cast
+  %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<*xf32>
+
+  return %1: tensor<*xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @linalg_effects(
 //  CHECK-SAME:     %[[A:[a-z0-9]*]]: tensor<?x?xf32>
 //  CHECK-SAME:     %[[B:[a-z0-9]*]]: memref<?x?xf32>
 //  CHECK-SAME:     %[[C:[a-z0-9]*]]: tensor<?x?xf32>
 func.func @linalg_effects(%a : tensor<?x?xf32>, %b : memref<?x?xf32>, %c : tensor<?x?xf32>) {
   // CHECK-NOT:   %{{.*}} = linalg.matmul
   %t = linalg.matmul ins(%a, %b : tensor<?x?xf32>, memref<?x?xf32>)
                     outs(%c : tensor<?x?xf32>) -> tensor<?x?xf32>