[Mlir-commits] [mlir] [MLIR][Tensor] Enhance bufferization of tensor.expand_shape op (PR #128871)

[llvmlistbot at llvm.org]

================
@@ -337,14 +337,27 @@ struct ExpandShapeOpInterface
     if (failed(buffer))
       return failure();
 
-    // Memref result type is inferred by the builder based on reassociation
-    // indices and result shape.
-    // TODO: Instead of inferring the output shape argument of
-    // memref.expand_shape op, use output_shape argument of tensor.expand_shape
-    // op.
-    replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
-        rewriter, op, tensorResultType.getShape(), *buffer,
-        expandShapeOp.getReassociationIndices());
+    // Use output_shape argument of tensor.expand_shape op to get the result
+    // shapes of the memref.expand_shape op to be created.
+    SmallVector<OpFoldResult> outShape;
----------------
MaheshRavishankar wrote:

I think you should be able to use `expandShapeOp.getMixedOutputShape()`.

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