[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 updated https://github.com/llvm/llvm-project/pull/217714

>From 970b956842b79f6963530fb6721c352ab5882046 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> {