[mlir][tensor][bufferize] Support 0-d collapse_shape with offset

[bkramer]

looks good, thanks

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     RankedTensorType tensorResultType = collapseShapeOp.getResultType();
     Value buffer =
         *state.getBuffer(rewriter, collapseShapeOp->getOpOperand(0) /*src*/);
 
     if (tensorResultType.getRank() == 0) {
       // 0-d collapses must go through a different op builder.
       auto bufferType = buffer.getType().cast<MemRefType>();
-      // Assume identity layout: No offset.
-      assert(bufferType.getLayout().isIdentity() &&
-             "non-zero offset for 0-d collapse not supported");
+      MemRefType resultType;
+
+      if (bufferType.getLayout().isIdentity()) {
+        // Standard layout: result type has no offset.
         MemRefLayoutAttrInterface layout;
-      auto resultType = MemRefType::get({}, tensorResultType.getElementType(),
-                                        layout, bufferType.getMemorySpace());
+        resultType = MemRefType::get({}, tensorResultType.getElementType(),
+      } else {
+        // Source memref has a layout map: result type has the same offset as
+        // the source type.
+        SmallVector<int64_t> strides;
+        int64_t offset;
+        if (failed(getStridesAndOffset(bufferType, strides, offset)))
+          return failure();
+        AffineMap resultLayout =
+            makeStridedLinearLayoutMap({}, offset, op->getContext());
+        resultType =
+            MemRefType::get({}, tensorResultType.getElementType(), resultLayout,
+                            bufferType.getMemorySpaceAsInt());
+      }
+
       replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
           rewriter, op, resultType, buffer, collapseShapeOp.reassociation());
       return success();
     }
 
     // Result type is inferred by the builder.
     replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
         rewriter, op, buffer, collapseShapeOp.getReassociationIndices());

mlir/test/Dialect/Tensor/bufferize.mlir

 // RUN: mlir-opt %s -tensor-bufferize | FileCheck %s
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 140 + d1 * 20 + d2 * 5 + d3 + s0)>
+// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0 + 1)>
+// CHECK-DAG: #[[$MAP4:.*]] = affine_map<() -> (1)>
 
 // CHECK-LABEL:   func @dim(
 // CHECK-SAME:              %[[TENSOR:.*]]: tensor<f32>,
 // CHECK-SAME:              %[[INDEX:.*]]: index) -> index {
 // CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<f32>
 // CHECK:           %[[EXTENT:.*]] = memref.dim %[[MEMREF]], %[[INDEX]] : memref<f32>
 // CHECK:           return %[[EXTENT]] : index
 func @dim(%arg0: tensor<f32>, %arg1: index) -> index {
func @tensor.collapse_shape_to_scalar(%t1: tensor<1x1x1xf32>) -> tensor<f32> {
   // CHECK: %[[collapsed:.*]] = memref.collapse_shape %[[m1]] [] : memref<1x1x1xf32> into memref<f32>
   %0 = tensor.collapse_shape %t1 []
       : tensor<1x1x1xf32> into tensor<f32>
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[collapsed]]
   // CHECK: return %[[r]]
   return %0 : tensor<f32>
 }
+
+// CHECK-LABEL: func @tensor.collapse_shape_of_slice(
+func @tensor.collapse_shape_of_slice(%arg0: tensor<2xi32>) -> tensor<i32> {
+  // CHECK: memref.subview %{{.*}}[1] [1] [1] : memref<2xi32> to memref<1xi32, #[[$MAP3]]>
+  %0 = tensor.extract_slice %arg0[1] [1] [1] : tensor<2xi32> to tensor<1xi32>
+  // CHECK: memref.collapse_shape %{{.*}} [] : memref<1xi32, #[[$MAP3]]> into memref<i32, #[[$MAP4]]>
+  %1 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  return %1 : tensor<i32>
+}