Tensor Dialect canonicalizer `FoldTensorCastProducerOp` can produce invalid IR

[christopherbate]

Reproducer:

// RUN: mlir-opt -canonicalize
func.func @materialize_in_destination_tensor_cast(%arg0: tensor<4xf32>, %arg1: index) -> tensor<?xf32> {
  %0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32>
  %1 = tensor.cast %arg0 : tensor<4xf32> to tensor<?xf32>
  %2 = bufferization.materialize_in_destination %1 in %0 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
  return %2 : tensor<?xf32>
}

FoldTensorCastProducerOp defined in lib/Dialect/Tensor/IR/TensorOps.cpp assumes that any operand of an op that implements DestinationStyleOpInterface can absorb a tensor.cast operation that is erasing information from the type. However, that assumption is not encoded anywhere in the spec of DestinationStyleOpInterface. For example, bufferization.materialize_in_destination` requires that the shapes of its DPS input and DPS output operands match, and therefore the canonicalizer will produce IR that fails verification :

within split at test/Dialect/Bufferization/canonicalize.mlir:392 offset :6:8: error: 'bufferization.materialize_in_destination' op failed to verify that all of {source, dest} have same shape
  %2 = bufferization.materialize_in_destination %1 in %0 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
       ^
within split at test/Dialect/Bufferization/canonicalize.mlir:392 offset :6:8: note: see current operation: %1 = "bufferization.materialize_in_destination"(%arg0, %0) : (tensor<4xf32>, tensor<?xf32>) -> tensor<?xf32>

[matthias-springer]

#92681 fixes the issue for materialize_in_destination. A more general fix would make the canonicalization pattern opt-in for each op that supports it. There are already a few ops that are blacklisted in the canonicalization pattern (e.g., scf.forall) and it is not always possible/practical to add additional ops to that list.

(Generalizing materialize_in_destination as in #92681 still makes sense because memref.copy also supports mixed static/dynamic dimensions.)

[joker-eph]

This looks more like a workaround than a fix to me: a canonicalized shouldn't change the type of an SSA value without more information

[matthias-springer]

I agree, it's likely still broken for other ops. I tried to make this pattern opt-in for each op that supports it. I.e., only ops that support the pattern should add it in their getCanonicalizationPatterns function; and rooting the pattern only the specified op name.

This requires a few more cleanups first, so I didn't get to it yet. E.g., tensor.pack, currently uses hasCanonicalizeMethod but we want hasCanonicalizer so that I can add the pattern. (I think the two cannot be used at the same time?)

[christopherbate]

I agree, it's likely still broken for other ops. I tried to make this pattern opt-in for each op that supports it. I.e., only ops that support the pattern should add it in their getCanonicalizationPatterns function; and rooting the pattern only the specified op name.

@matthias-springer @joker-eph I opened PR #91382 two weeks ago to do just this. That PR makes the pattern opt-in.