[Mlir-commits] [mlir] [mlir][tensor] Fix FoldTensorCastProducerOp for multiple result operations (PR #93374)

[Han-Chung Wang]

================
@@ -4531,17 +4531,18 @@ struct FoldTensorCastProducerOp
     if (!hasTensorCastOperand)
       return failure();
 
-    SmallVector<Type, 4> newResultTypes;
-    newResultTypes.reserve(op->getNumResults());
+    SmallVector<Type, 4> newResultTypes(op->getResultTypes());
     SmallVector<Value, 4> newOperands;
     newOperands.reserve(op->getNumOperands());
+    // Assumes that the result has dpsInits followed by nonDpsInits.
----------------
hanhanW wrote:

I had the same confusion. Let me explain a bit more..  The number of results is not as same as the number of init operands. It is not a requirement of DestinationPassingStyleInterface. In this example, it has one init tensor and two result types (i.e., tensor + i32 scalar).

We were confused that the mapping between result types and init tensor. The code is wrong if we have `i32, tensor<xxx>` return types. After reading the doc again, I now think that the assumption is correct. The leading result types should match init tensor types. 

```
    Each tensor init operand is tied to a corresponding tensor OpResult in a
    1-to-1 fashion. The i-th init tensor is tied to the i-th OpResult. The op
    may not have any additional OpResults. Init operands and their tied
    OpResults have the same type.
```

(It is not verified in the implementation, so I thought that it's not documented.)

https://github.com/llvm/llvm-project/blob/7476c20c481cbccbdb89139fb94620e083015932/mlir/lib/Interfaces/DestinationStyleOpInterface.cpp#L29-L62

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