[mlir][linalg][bufferize] Support TensorExpand/CollapseShape

[springerm]

Summary

Depends On D112511

[springerm]

update

[springerm]

rebase

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

1993 ↗

(On Diff #382495)

can we use something like:

llvm::map_iterator(
  attr.cast<ArrayAttr>().getAsValueRange<IntegerAttr>(),
  [](const APInt &val) { return !val.getInt(); }))

?

1997 ↗

(On Diff #382495)

dyn_cast + assert is an anti-pattern FWIH, you can just cast

[springerm]

update

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

assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
            "CallOp::bufferize: cast incompatible");
     replaceOpWithNewBufferizedOp<memref::CastOp>(rewriter, op, resultMemRefType,
                                                  *resultBuffer);
 
     return success();
   }
 };
 
+/// Bufferization of tensor.collapse_shape. Replace with memref.collapse_shape.
+struct CollapseShapeOpInterface
+    : public BufferizableOpInterface::ExternalModel<CollapseShapeOpInterface,
+                                                    tensor::CollapseShapeOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const BufferizationState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult>
+  getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                      const BufferizationState &state) const {
+    if (&opOperand == &op->getOpOperand(0) /*src*/)
+      return {op->getOpResult(0)};
+    return {};
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const BufferizationState &state) const {
+    return BufferRelation::Equivalent;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationState &state) const {
+    auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
+    Value buffer =
+        *state.getBuffer(rewriter, collapseShapeOp->getOpOperand(0) /*src*/);
+    Type resultType =
+        getMemRefType(collapseShapeOp.getResultType(), state.getOptions());
+    replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
+        rewriter, op, resultType, buffer, collapseShapeOp.reassociation());
+    return success();
+  }
+};
+
 /// Bufferization of tensor.dim. Replace with memref.dim.
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const BufferizationState &state) const {
     return true;
   }
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
     Value v = *state.getBuffer(rewriter, dimOp->getOpOperand(0) /*source*/);
     replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, v, dimOp.index());
     return success();
   }
 };
 
+/// Bufferization of tensor.expand_shape. Replace with memref.expand_shape.
+struct ExpandShapeOpInterface
+    : public BufferizableOpInterface::ExternalModel<ExpandShapeOpInterface,
+                                                    tensor::ExpandShapeOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const BufferizationState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult>
+  getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                      const BufferizationState &state) const {
+    if (&opOperand == &op->getOpOperand(0) /*src*/)
+      return {op->getOpResult(0)};
+    return {};
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const BufferizationState &state) const {
+    return BufferRelation::Equivalent;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationState &state) const {
+    auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
+    Value buffer =
+        *state.getBuffer(rewriter, expandShapeOp->getOpOperand(0) /*src*/);
+    Type resultType =
+        getMemRefType(expandShapeOp.getResultType(), state.getOptions());
+    replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
+        rewriter, op, resultType, buffer, expandShapeOp.reassociation());
+    return success();
+  }
+};
+
 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
 struct ExtractSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
                                                     tensor::ExtractSliceOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const BufferizationState &state) const {
     return false;
   }

 
 } // namespace
 } // namespace tensor
 } // namespace mlir
 
 void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<CastOp, CastOpInterface>();
+  registry.addOpInterface<CollapseShapeOp, CollapseShapeOpInterface>();
   registry.addOpInterface<DimOp, DimOpInterface>();
+  registry.addOpInterface<ExpandShapeOp, ExpandShapeOpInterface>();
   registry.addOpInterface<ExtractSliceOp, ExtractSliceOpInterface>();
   registry.addOpInterface<ExtractOp, ExtractOpInterface>();
   registry.addOpInterface<FromElementsOp, FromElementsOpInterface>();
   registry.addOpInterface<GenerateOp, GenerateOpInterface>();
   registry.addOpInterface<InsertOp, InsertOpInterface>();
   registry.addOpInterface<InsertSliceOp, InsertSliceOpInterface>();
   registry.addOpInterface<RankOp, RankOpInterface>();
 }

mlir/test/Dialect/Tensor/bufferize.mlir

func @tensor.insert(%t1: tensor<5xf32>, %idx1: index, %f: f32) -> tensor<5xf32> {
   // CHECK: memref.copy %[[m1]], %[[alloc]]
   // CHECK: memref.store %[[f]], %[[alloc]][%[[idx1]]]
   %0 = tensor.insert %f into %t1[%idx1] : tensor<5xf32>
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK: return %[[r]]
   return %0 : tensor<5xf32>
 }
+
+// CHECK-LABEL: func @tensor.expand_shape(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?x10xf32>
+func @tensor.expand_shape(%t1: tensor<?x10xf32>) -> tensor<2x?x10xf32> {
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xf32>
+  // CHECK: %[[expanded:.*]] = memref.expand_shape %[[m1]] [
+  // CHECK-SAME: [0, 1], [2]] : memref<?x10xf32> into memref<2x?x10xf32>
+  %0 = tensor.expand_shape %t1 [[0, 1], [2]]
+      : tensor<?x10xf32> into tensor<2x?x10xf32>
+
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
+  // CHECK: return %[[r]]
+  return %0 : tensor<2x?x10xf32>
+}
+
+// CHECK-LABEL: func @tensor.collapse_shape(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<2x?x?xf32>
+func @tensor.collapse_shape(%t1: tensor<2x?x?xf32>) -> tensor<?x?xf32> {
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<2x?x?xf32>
+  // CHECK: %[[collapsed:.*]] = memref.collapse_shape %[[m1]] [
+  // CHECK-SAME: [0, 1], [2]] : memref<2x?x?xf32> into memref<?x?xf32>
+  %0 = tensor.collapse_shape %t1 [[0, 1], [2]]
+      : tensor<2x?x?xf32> into tensor<?x?xf32>
+
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[collapsed]]
+  // CHECK: return %[[r]]
+  return %0 : tensor<?x?xf32>
+}