[mlir][Linalg] Fix error in `LinalgOp` -> `tensor::CastOp` folding.

[mravishankar]

Summary

When replacing the uses of the LinalgOp other than the use in the
tensor.cast op, check if the op itself has multiple uses, instead of
the specific result having multiple uses.

[gysit]

I think https://reviews.llvm.org/D121369 already fixes the problem and additionally makes the pattern work with the transform dialect. Landing the test still make sense!

[mravishankar]

In D121648#3381598, @gysit wrote:

I think https://reviews.llvm.org/D121369 already fixes the problem and additionally makes the pattern work with the transform dialect. Landing the test still make sense!

Ah I see this is the same commit you pointed me to. Didnt connect the two. Ill land the test then.

[mravishankar]

Dropping cod echange and only keeping the test.

mlir/test/Dialect/Linalg/canonicalize.mlir

 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>)
 //       CHECK:  %[[OUT_CAST:.+]] = tensor.cast %[[ARG2]] : tensor<?x?x?x?xf32> to tensor<4x8x12x16xf32>
 //       CHECK:  %[[CONV:.+]] = linalg.conv_2d_nchw_fchw
 //  CHECK-SAME:      ins(%[[ARG0]], %[[ARG1]] :
 //  CHECK-SAME:      outs(%[[OUT_CAST]] :
 //       CHECK:  %[[RESULT_CAST:.+]] = tensor.cast %[[CONV]]
 //       CHECK:  return %[[CONV]], %[[RESULT_CAST]]
+
+// -----
+
+func @fold_multi_use_generic_op_with_consumer(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<2x3x4xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %init1 = linalg.init_tensor [%d1, %d2, %d0] : tensor<?x?x?xf32>
+  %init2 = linalg.init_tensor [%d2, %d1, %d0] : tensor<?x?x?xf32>
+  %0:2 = linalg.generic {
+      iterator_types = ["parallel", "parallel", "parallel"],
+      indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                       affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
+                       affine_map<(d0, d1, d2) -> (d2, d1, d0)>]}
+      ins(%arg0 : tensor<?x?x?xf32>) outs(%init1, %init2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32) :
+      linalg.yield %b0, %b0 : f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  %1 = tensor.cast %0#1 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %0#0, %1 : tensor<?x?x?xf32>, tensor<2x3x4xf32>
+}
+//       CHECK: func @fold_multi_use_generic_op_with_consumer
+//  CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>
+//   CHECK-DAG:   %[[INIT1:.+]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
+//   CHECK-DAG:   %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x3x2xf32>
+//   CHECK-DAG:   %[[INIT2:.+]] = linalg.init_tensor [3, 2, 4] : tensor<3x2x4xf32>
+//       CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
+//  CHECK-SAME:       ins(%[[CAST]] :
+//  CHECK-SAME:       outs(%[[INIT2]], %[[INIT1]] :
+//       CHECK:   %[[RETURN_CAST:.+]] = tensor.cast %[[GENERIC]]#0 : tensor<3x2x4xf32> to tensor<?x?x?xf32>
+//       CHECK:   return %[[RETURN_CAST]], %[[GENERIC]]#1