llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Victor Perez (victor-eds)
<details>
<summary>Changes</summary>
`ExpandShapeOpInterface::bufferize` creates a `memref.expand_shape` directly without consulting `getBufferType`. The `memref::ExpandShapeOp` builder computes the result layout with `computeExpandedType` and asserts when that fails. `computeExpandedType` fails when the source layout is not identity and `getStridesAndOffset` cannot decompose it, so a non-strided affine layout aborts the compiler instead of reporting a failed bufferization.
Query `getBufferType` first and return failure when it fails. The model's own `getBufferType` already propagates the `computeExpandedType` failure.
Rest of the behaviour is unchanged: we just fail gracefully now instead of crashing.
---
Code authored by Claude Code.
---
Full diff: https://github.com/llvm/llvm-project/pull/217714.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp (+5-2)
- (modified) mlir/test/Dialect/Tensor/bufferize.mlir (+35-1)
``````````diff
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 525d57341cee5..fa02d23272503 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -343,14 +343,17 @@ struct ExpandShapeOpInterface
const BufferizationOptions &options,
BufferizationState &state) const {
auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
- auto tensorResultType = expandShapeOp.getResultType();
+ FailureOr<BufferLikeType> maybeResultType =
+ bufferization::getBufferType(expandShapeOp.getResult(), options, state);
+ if (failed(maybeResultType))
+ return failure();
FailureOr<Value> buffer =
getBuffer(rewriter, expandShapeOp.getSrc(), options, state);
if (failed(buffer))
return failure();
auto memrefExpandShape = memref::ExpandShapeOp::create(
- rewriter, op->getLoc(), tensorResultType.getShape(), *buffer,
+ rewriter, op->getLoc(), *maybeResultType, *buffer,
expandShapeOp.getReassociationIndices(),
expandShapeOp.getMixedOutputShape());
replaceOpWithBufferizedValues(rewriter, op,
diff --git a/mlir/test/Dialect/Tensor/bufferize.mlir b/mlir/test/Dialect/Tensor/bufferize.mlir
index 334bfa8111e39..965742aacc579 100644
--- a/mlir/test/Dialect/Tensor/bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/bufferize.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --one-shot-bufferize="dialect-filter=tensor,bufferization copy-before-write unknown-type-conversion=identity-layout-map" -cse -split-input-file | FileCheck %s
+// RUN: mlir-opt %s --one-shot-bufferize="dialect-filter=tensor,bufferization copy-before-write unknown-type-conversion=identity-layout-map" -cse -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @dim(
// CHECK-SAME: %[[TENSOR:.*]]: tensor<*xf32>,
@@ -426,6 +426,40 @@ func.func @tensor.expand_shape_multiple_dynamic_indices(%t1: tensor<?x256xf32>,
}
// -----
+// CHECK-LABEL: func @tensor.expand_shape_of_strided_layout(
+// CHECK-SAME: %[[m1:.*]]: memref<4x4xf32, strided<[8, 1]>>
+func.func @tensor.expand_shape_of_strided_layout(
+ %m1: memref<4x4xf32, strided<[8, 1]>>) -> tensor<2x2x4xf32> {
+ %t1 = bufferization.to_tensor %m1 restrict
+ : memref<4x4xf32, strided<[8, 1]>> to tensor<4x4xf32>
+
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.copy
+ // CHECK: %[[expanded:.*]] = memref.expand_shape %[[m1]] {{\[\[}}0, 1], [2]] output_shape [2, 2, 4] : memref<4x4xf32, strided<[8, 1]>> into memref<2x2x4xf32, strided<[16, 8, 1]>>
+ %0 = tensor.expand_shape %t1 [[0, 1], [2]] output_shape [2, 2, 4]
+ : tensor<4x4xf32> into tensor<2x2x4xf32>
+
+ // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
+ // CHECK: return %[[r]]
+ return %0 : tensor<2x2x4xf32>
+}
+// -----
+
+#nonstrided = affine_map<(d0, d1) -> (d0 * 4 + d1 floordiv 2)>
+
+func.func @tensor.expand_shape_of_non_strided_layout(
+ %m1: memref<4x4xf32, #nonstrided>) -> tensor<2x2x4xf32> {
+ %t1 = bufferization.to_tensor %m1 restrict
+ : memref<4x4xf32, #nonstrided> to tensor<4x4xf32>
+
+ // expected-error @below{{failed to bufferize op}}
+ %0 = tensor.expand_shape %t1 [[0, 1], [2]] output_shape [2, 2, 4]
+ : tensor<4x4xf32> into tensor<2x2x4xf32>
+
+ return %0 : tensor<2x2x4xf32>
+}
+// -----
+
// CHECK-LABEL: func @tensor.collapse_shape(
// CHECK-SAME: %[[t1:.*]]: tensor<2x?x?xf32>
func.func @tensor.collapse_shape(%t1: tensor<2x?x?xf32>) -> tensor<?x?xf32> {
``````````
</details>
https://github.com/llvm/llvm-project/pull/217714