[mlir][sparse][bufferize] Implement BufferizableOpInterface

[springerm]

Summary

Only the analysis part of the interface is implemented. The bufferization itself is performed by the SparseTensorConversion pass.

Depends On D128137

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h

2

This is not really the Impl, right, since this is header. You have the same comment in the CPP file.

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

23

do we need those using, given that you add the namespace and fully qualify outside of it?

36

shall we remove the folding, just to keep things simple, and make that the "copy" op?

[springerm]

rebase

[springerm]

mlir/include/mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h

2

Naming is a bit confusing, but this Impl refers to "Interface Implementation", not source vs header. registerBufferizableOpInterfaceExternalModels is a function that registers the impls.

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

36

AllocTensorOp is how we copy during bufferization. It doesn't really matter (for bufferization purposes) if this folds away or not. We anyway have to return true here because sparse -> dense or the other way around definitely allocates.

[springerm]

address comments

[aartbik]

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

15

Note that, as discussed, this will actually fail verification after the upcoming change.

[aartbik]

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

15

You can avoid that by adding

%1 = sparse_tensor.load %0 : tensor<20x40xf32, SDCSR>
return %1 : tensor<20x40xf32, SDCSR>

note that it still escapes, but just not directly (i.e. not unitialized, since the load op wraps-up the "insertion").

[springerm]

update

mlir/include/mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h

+//===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H_
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H_
+
+namespace mlir {
+class DialectRegistry;
+
+namespace sparse_tensor {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H_

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
inline void registerAllDialects(DialectRegistry &registry) {
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
+  sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

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

+//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// These BufferizableOpInterface implementations provide analysis-related
+// interface methods only. They are getting bufferized by the
+// SparseTensorConversion pass.
+
+#include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir::bufferization;
+using namespace mlir::sparse_tensor;
+
+namespace mlir {
+namespace sparse_tensor {
+namespace {
+
+struct ConvertOpInterface
+    : public BufferizableOpInterface::ExternalModel<ConvertOpInterface,
+                                                    sparse_tensor::ConvertOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    // ConvertOps may allocate. (Unless they convert between two identical
+    // types, then they fold away.)
+    return true;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+
+  bool isWritable(Operation *op, Value value,
+                  const AnalysisState &state) const {
+    return true;
+  }
+};
+
+struct LoadOpInterface
+    : public BufferizableOpInterface::ExternalModel<LoadOpInterface,
+                                                    sparse_tensor::LoadOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {op->getOpResult(0)};
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const AnalysisState &state) const {
+    return BufferRelation::Equivalent;
+  }
+};
+
+struct NewOpInterface
+    : public BufferizableOpInterface::ExternalModel<NewOpInterface,
+                                                    sparse_tensor::NewOp> {
+  bool isMemoryWrite(Operation *op, OpResult opResult,
+                     const AnalysisState &state) const {
+    // NewOps allocate but do not write.
+    return false;
+  }
+
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+};
+
+struct ReleaseOpInterface
+    : public BufferizableOpInterface::ExternalModel<ReleaseOpInterface,
+                                                    sparse_tensor::ReleaseOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
+} // namespace
+} // namespace sparse_tensor
+} // namespace mlir
+
+void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(
+      +[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) {
+        sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
+        sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
+        sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
+        sparse_tensor::ReleaseOp::attachInterface<ReleaseOpInterface>(*ctx);
+      });
+}

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorTransforms
+  BufferizableOpInterfaceImpl.cpp
   CodegenUtils.cpp
   Sparsification.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmeticDialect
   MLIRBufferizationDialect
+  MLIRBufferizationTransforms
   MLIRComplexDialect
   MLIRFuncDialect
   MLIRIR
   MLIRLLVMDialect
   MLIRLinalgDialect
   MLIRLinalgTransforms
   MLIRMemRefDialect
   MLIRPass

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

+// RUN: mlir-opt %s -tensor-copy-insertion="allow-return-allocs" | FileCheck %s
+// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" | FileCheck %s --check-prefix=CHECK-FUNC
+
+#DCSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (i,j)>
+}>
+
+// CHECK-LABEL: func @bufferization_alloc_tensor
+// CHECK-FUNC-LABEL: func @bufferization_alloc_tensor
+func.func @bufferization_alloc_tensor() -> tensor<20x40xf32, #DCSR> {
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]}
+  // CHECK-FUNC: bufferization.alloc_tensor() {bufferization.escape = [true]}
+  %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
+  %1 = sparse_tensor.load %0 : tensor<20x40xf32, #DCSR>
+  return %1 : tensor<20x40xf32, #DCSR>
+}
+
+!Filename = !llvm.ptr<i8>
+// CHECK-LABEL: func @sparse_tensor_new
+// CHECK-FUNC-LABEL: func @sparse_tensor_new
+func.func @sparse_tensor_new(%file: !Filename) -> tensor<20x40xf32, #DCSR> {
+  // CHECK: sparse_tensor.new {{.*}} {bufferization.escape = [false]}
+  // CHECK-FUNC: sparse_tensor.new {{.*}} {bufferization.escape = [true]}
+  %0 = sparse_tensor.new %file : !Filename to tensor<20x40xf32, #DCSR>
+  return %0 : tensor<20x40xf32, #DCSR>
+}
+
+// CHECK-LABEL: func @sparse_tensor_convert
+// CHECK-FUNC-LABEL: func @sparse_tensor_convert
+func.func @sparse_tensor_convert() -> tensor<20x40xf32> {
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false]}
+  // CHECK-FUNC: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false]}
+  %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
+  // CHECK: %[[loaded:.*]] = sparse_tensor.load %[[alloc]]
+  // CHECK-FUNC: %[[loaded:.*]] = sparse_tensor.load %[[alloc]]
+  %1 = sparse_tensor.load %0 : tensor<20x40xf32, #DCSR>
+  // CHECK: sparse_tensor.convert %[[loaded]] {bufferization.escape = [false]}
+  // CHECK-FUNC: sparse_tensor.convert %[[loaded]] {bufferization.escape = [true]}
+  %2 = sparse_tensor.convert %1 : tensor<20x40xf32, #DCSR> to tensor<20x40xf32>
+  return %2 : tensor<20x40xf32>
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob([
         "lib/Dialect/SparseTensor/Transforms/*.cpp",
         "lib/Dialect/SparseTensor/Transforms/*.h",
     ]),
     hdrs = [
+        "include/mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h",
         "include/mlir/Dialect/SparseTensor/Transforms/Passes.h",
         "include/mlir/ExecutionEngine/SparseTensorUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithmeticDialect",
         ":BufferizationDialect",
+        ":BufferizationTransforms",
         ":ComplexDialect",
         ":FuncDialect",
         ":FuncTransforms",
         ":IR",
         ":LLVMDialect",
         ":LinalgDialect",
         ":LinalgTransforms",
         ":LinalgUtils",