[mlir][tensor] Realloc to bufferize ReshapeOp with non-identity layout.

[ingomueller-net]

Summary

tensor.reshape gets bufferized to memref.reshape, which requires its
operands to have statically known identity layout. The previous
implementation did this conversion blindly, so a tensor in non-identity
layout would be converted into a memref op that does not verify. This
patch extends the conversion such that an input that gets bufferized
into a non-identity layout gets copied into a new allocation, which can
then be reshaped.

[ingomueller-net]

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

1006

I do not intend to work on this TODO as part of this patch.

[ingomueller-net]

@springerm: Is this patch completely obsolete now? Should we abandon it?

[springerm]

Something like this has landed recently: https://github.com/llvm/llvm-project/pull/65654

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     return success();
   }
 };
 
 /// Bufferization of tensor.reshape. Replace with memref.reshape.
 struct ReshapeOpInterface
     : public BufferizableOpInterface::ExternalModel<ReshapeOpInterface,
                                                     tensor::ReshapeOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     if (&opOperand == &op->getOpOperand(1) /* shape */)
       return true;
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
   AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto reshapeOp = cast<tensor::ReshapeOp>(op);
+    Location loc = op->getLoc();
+
+    // Get bufferized src and shape operands.
     FailureOr<Value> srcBuffer =
         getBuffer(rewriter, reshapeOp.getSource(), options);
     FailureOr<Value> shapeBuffer =
         getBuffer(rewriter, reshapeOp.getShape(), options);
     if (failed(srcBuffer) || failed(shapeBuffer))
       return failure();
+
+    // Copy src buffer into new allocation if it doesn't have identity layout.
+    // TODO: This could be extended to layouts that are not the identity layout
+    //       but still allow zero-copy reshape (such as a 1D input with an
+    //       offset).
+    auto bufferType = cast<MemRefType>(srcBuffer->getType());
+    if (!bufferType.getLayout().isIdentity()) {
+      bool dealloc = shouldDeallocateOpResult(
+          cast<OpResult>(reshapeOp.getResult()), options);
+      FailureOr<Value> tensorAlloc =
+          allocateTensorForShapedValue(rewriter, loc, reshapeOp.getSource(),
+                                       /*escape=*/!dealloc, options,
+                                       /*copy=*/true);
+      if (failed(tensorAlloc))
+        return failure();
+      auto memrefType =
+          MemRefType::get(bufferType.getShape(), bufferType.getElementType(),
+                          AffineMap(), bufferType.getMemorySpace());
+      auto toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
+          op->getLoc(), memrefType, *tensorAlloc);
+      srcBuffer = toMemrefOp.getResult();
+    }
+
+    // Convert this op to memref equivalent on buffers.
     auto maybeResultMemRefType =
         bufferization::getBufferType(reshapeOp.getResult(), options);
     if (failed(maybeResultMemRefType))
       return failure();
     replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
         rewriter, op, maybeResultMemRefType.value(), *srcBuffer, *shapeBuffer);
     return success();
   }

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

func.func @tensor.reshape() -> tensor<2x2x5xf32> {
   %shape = arith.constant dense<[2, 2, 5]> : tensor<3xi64>
 
   // CHECK: %[[RESHAPED:.*]] = memref.reshape %[[M1]](%[[SHAPE]]) : (memref<?x10xf32>, memref<3xi64>) -> memref<2x2x5xf32>
   %reshaped = tensor.reshape %t1(%shape) : (tensor<?x10xf32>, tensor<3xi64>) -> tensor<2x2x5xf32>
 
   // CHECK: return %[[RESHAPED]]
   return %reshaped : tensor<2x2x5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor.reshape_non_identity(
+//  CHECK-SAME:     %[[ARG0:.*]]: memref<?x10xf32, strided<[?, ?], offset: ?>>) -> memref<2x2x5xf32> {
+func.func @tensor.reshape_non_identity(%t1 : tensor<?x10xf32>) -> tensor<2x2x5xf32> {
+  // CHECK: %[[SHAPE:.*]] = memref.get_global @{{.*}} : memref<3xi64>
+  %shape = arith.constant dense<[2, 2, 5]> : tensor<3xi64>
+
+  // CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) {alignment = 64 : i64} : memref<?x10xf32>
+  // CHECK: memref.copy %[[ARG0]], %[[ALLOC]] : memref<?x10xf32, strided<[?, ?], offset: ?>> to memref<?x10xf32>
+  // CHECK: %[[RESHAPED:.*]] = memref.reshape %[[ALLOC]](%[[SHAPE]]) : (memref<?x10xf32>, memref<3xi64>) -> memref<2x2x5xf32>
+  %reshaped = tensor.reshape %t1(%shape) : (tensor<?x10xf32>, tensor<3xi64>) -> tensor<2x2x5xf32>
+
+  // CHECK: return %[[RESHAPED]]
+  return %reshaped : tensor<2x2x5xf32>
+}