[mlir][linalg] Fix `FoldTensorCastProducerOp` for generic with memref output

[Hardcode84]

Summary

Type should only be added to results if it is tensor.

[pifon2a]

Just out of curiosity. Why do you need the mixed form of linalg.generic?

[Hardcode84]

Just out of curiosity. Why do you need the mixed form of linalg.generic?

In our python compiler we lowering numpy to the mix of tensor and memref (we cannot lower everything to tensor land with the presence on like ops like setitem)

[Hardcode84]

ping

[nicolasvasilache]

In D137801#3921101, @Hardcode84 wrote:

Just out of curiosity. Why do you need the mixed form of linalg.generic?

In our python compiler we lowering numpy to the mix of tensor and memref (we cannot lower everything to tensor land with the presence on like ops like setitem)

This is very interesting and a use case we have always wanted to support despite our lacking a concrete example.
We speculated that bufferization could want to be done partially but in reality this need never materialized.

Some short post describing the use case a bit more and stressing the need for ops that work with both tensors and buffers would be quite useful IMO :)

I would also go as far as adding a comment in the @cast_producer_mixed test that mentions some of this, lest you don't mind risking someone in the future coming and saying "oh the mixed form is only used in that 1 test, let's drop it.." :)

Thanks for pushing on this!

[Hardcode84]

In D137801#3931492, @nicolasvasilache wrote:

In D137801#3921101, @Hardcode84 wrote:

Just out of curiosity. Why do you need the mixed form of linalg.generic?

In our python compiler we lowering numpy to the mix of tensor and memref (we cannot lower everything to tensor land with the presence on like ops like setitem)

This is very interesting and a use case we have always wanted to support despite our lacking a concrete example.
We speculated that bufferization could want to be done partially but in reality this need never materialized.

Some short post describing the use case a bit more and stressing the need for ops that work with both tensors and buffers would be quite useful IMO :)

I would also go as far as adding a comment in the @cast_producer_mixed test that mentions some of this, lest you don't mind risking someone in the future coming and saying "oh the mixed form is only used in that 1 test, let's drop it.." :)

Thanks for pushing on this!

We are planning to present our numpy dialect (https://github.com/intel/mlir-extensions/blob/main/mlir/include/imex/Dialect/ntensor/IR/NTensorOps.td) and lowering on the MLIR ODM eventually, but we need more time to prepare. Specifically for the mixed generic's , we are using them to lower numpy calsl with out parameters. We lower entire numpy function to linalg-on-tensors and then insert final linalg.generic which copies result to out memref (https://github.com/intel/mlir-extensions/blob/main/mlir/lib/Conversion/NtensorToLinalg.cpp#L69), and hope it all will be fused together.

[nicolasvasilache]

In D137801#3931536, @Hardcode84 wrote:

In D137801#3931492, @nicolasvasilache wrote:

In D137801#3921101, @Hardcode84 wrote:

Just out of curiosity. Why do you need the mixed form of linalg.generic?

In our python compiler we lowering numpy to the mix of tensor and memref (we cannot lower everything to tensor land with the presence on like ops like setitem)

This is very interesting and a use case we have always wanted to support despite our lacking a concrete example.
We speculated that bufferization could want to be done partially but in reality this need never materialized.

Some short post describing the use case a bit more and stressing the need for ops that work with both tensors and buffers would be quite useful IMO :)

I would also go as far as adding a comment in the @cast_producer_mixed test that mentions some of this, lest you don't mind risking someone in the future coming and saying "oh the mixed form is only used in that 1 test, let's drop it.." :)

Thanks for pushing on this!

We are planning to present our numpy dialect (https://github.com/intel/mlir-extensions/blob/main/mlir/include/imex/Dialect/ntensor/IR/NTensorOps.td) and lowering on the MLIR ODM eventually, but we need more time to prepare. Specifically for the mixed generic's , we are using them to lower numpy calsl with out parameters. We lower entire numpy function to linalg-on-tensors and then insert final linalg.generic which copies result to out memref (https://github.com/intel/mlir-extensions/blob/main/mlir/lib/Conversion/NtensorToLinalg.cpp#L69), and hope it all will be fused together.

Very cool, thanks for sharing!

and hope it all will be fused together.

Well.. we better be sure we give you all the tools you need :)

[Hardcode84]

rebase, add comment

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

for (auto *input : op.getDpsInputOperands()) {
                                 ? tensorCastOp.getSource()
                                 : input->get());
     }
     // Init tensors may fold, in which case the resultType must also change.
     for (auto *output : op.getDpsInitOperands()) {
       auto tensorCastOp = output->get().getDefiningOp<tensor::CastOp>();
       bool fold = canFoldIntoConsumerOp(tensorCastOp);
       newOperands.push_back(fold ? tensorCastOp.getOperand() : output->get());
+      if (!newOperands.back().getType().isa<MemRefType>())
         newResultTypes.push_back(newOperands.back().getType());
     }
     // Clone op.
     Operation *newOp =
         op.clone(rewriter, op->getLoc(), newResultTypes, newOperands);
     SmallVector<Value, 4> replacements;
     replacements.reserve(newOp->getNumResults());
     for (auto result : llvm::zip(op->getResults(), newOp->getResults())) {
       Value oldResult = std::get<0>(result);

mlir/test/Dialect/Linalg/canonicalize.mlir

func.func @dedeplicate_regression_test(%0: tensor<4xf32>, %1: memref<4xf32>) {
   %53 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>],
                         iterator_types = ["parallel"]}
                         outs(%36 : tensor<4xf32>) {
   ^bb0(%out: f32):
     linalg.yield %out : f32
   } -> tensor<4xf32>
   return
 }
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+func.func @cast_producer_mixed(%arg0 : tensor<5xf32>, %arg1: memref<?xf32>) {
+  %0 = tensor.cast %arg0 : tensor<5xf32> to tensor<?xf32>
+  linalg.generic {
+    indexing_maps = [#map, #map],
+    iterator_types = ["parallel"]
+  } ins(%0 : tensor<?xf32>)
+    outs(%arg1 : memref<?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32):
+    linalg.yield %arg2 : f32
+  }
+  return
+}
+
+// We need a mixed linalg as a bridge between tensor and memref worlds.
+// CHECK-LABEL: func @cast_producer_mixed
+//  CHECK-SAME:     (%[[ARG1:.*]]: tensor<5xf32>, %[[ARG2:.*]]: memref<?xf32>)
+//       CHECK:     linalg.generic {
+//  CHECK-SAME:    indexing_maps = [#map, #map],
+//  CHECK-SAME:    iterator_types = ["parallel"]
+//  CHECK-SAME:  } ins(%[[ARG1]] : tensor<5xf32>)
+//  CHECK-SAME:    outs(%[[ARG2]] : memref<?xf32>) {