[Mlir-commits] [mlir] [mlir][tensor] Fix assertion on bufferizing expand_shape with a non-strided layout (PR #217714)

[Victor Perez]

https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/217714

`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.

>From 51f718f84347fa295d4445e8a69c3a78b272858e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
 <victor.pc.upm at gmail.com>
Date: Thu, 20 Aug 2026 02:04:27 -0700
Subject: [PATCH] [mlir][tensor] Fix assertion on bufferizing expand_shape with
 a non-strided layout

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. Passing the
queried type to the builder also removes the asserting overload from the call
path.

Analysis behaviour is unchanged: bufferizesToMemoryRead, bufferizesToMemoryWrite
and getAliasingValues are untouched, so a strided source still bufferizes to a
pure view with no copy.
---
 .../BufferizableOpInterfaceImpl.cpp           |  7 ++--
 mlir/test/Dialect/Tensor/bufferize.mlir       | 36 ++++++++++++++++++-
 2 files changed, 40 insertions(+), 3 deletions(-)

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> {