[mlir][sparse] codegen for sparse dealloc

[bixia]

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

127

This change is not relevant to this PR?

[aartbik]

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

127

No, I noticed that we add the dialect *and* the ops. That is a bit redundant.

[aartbik]

rebased against main (mainly my own change)

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

 // to the SparseTensorConversion pass, eliminating the dependence on a runtime
 // support library, and providing much more opportunities for subsequent
 // compiler optimization of the generated code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
matchAndRewrite(tensor::CastOp op, OpAdaptor adaptor,
     auto encSrc = getSparseTensorEncoding(op.getSource().getType());
     if (!encDst || encDst != encSrc)
       return failure();
     rewriter.replaceOp(op, adaptor.getOperands());
     return success();
   }
 };
 
-/// Sparse conversion rule for pointer accesses.
+/// Sparse codegen rule for the dealloc operator.
+class SparseTensorDeallocConverter
+    : public OpConversionPattern<bufferization::DeallocTensorOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto enc = getSparseTensorEncoding(op.getTensor().getType());
+    if (!enc)
+      return failure();
+    // Replace the tuple deallocation with field deallocations.
+    Location loc = op->getLoc();
+    Value tuple = adaptor.getTensor();
+    for (unsigned i = 0, sz = tuple.getType().cast<TupleType>().size(); i < sz;
+         i++) {
+      Value mem = createTupleGet(rewriter, loc, tuple, i);
+      rewriter.create<memref::DeallocOp>(loc, mem);
+    }
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+/// Sparse codegen rule for pointer accesses.
 class SparseToPointersConverter : public OpConversionPattern<ToPointersOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToPointersOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Optional<int64_t> index = getConstantInt(adaptor.getOperands()[1]);
     if (!index)
       return failure();
     // Replace the requested pointer access with corresponding field.
     Location loc = op->getLoc();
     Value tuple = adaptor.getTensor();
     unsigned i = getFieldIndex(op.getTensor().getType(), /*ptrDim=*/*index, -1);
     rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
     return success();
   }
 };
 
-/// Sparse conversion rule for index accesses.
+/// Sparse codegen rule for index accesses.
 class SparseToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Optional<int64_t> index = getConstantInt(adaptor.getOperands()[1]);
     if (!index)
       return failure();
     // Replace the requested indices access with corresponding field.
     Location loc = op->getLoc();
     Value tuple = adaptor.getTensor();
     unsigned i = getFieldIndex(op.getTensor().getType(), -1, /*idxDim=*/*index);
     rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
     return success();
   }
 };
 
-/// Sparse conversion rule for value accesses.
+/// Sparse codegen rule for value accesses.
 class SparseToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Replace the requested values access with corresponding field.
     Location loc = op->getLoc();
     Value tuple = adaptor.getTensor();
-    unsigned i = tuple.getType().cast<TupleType>().size() - 1;  // last
+    unsigned i = tuple.getType().cast<TupleType>().size() - 1; // last
     rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//

 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseDimOpConverter, SparseCastConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToValuesConverter>(typeConverter, patterns.getContext());
+               SparseTensorDeallocConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToValuesConverter>(
+      typeConverter, patterns.getContext());
 }

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

target.addDynamicallyLegalOp<bufferization::AllocTensorOp>(
           return converter.isLegal(op.getType());
         });
     target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(
         [&](bufferization::DeallocTensorOp op) {
           return converter.isLegal(op.getTensor().getType());
         });
     // The following operations and dialects may be introduced by the
     // rewriting rules, and are therefore marked as legal.
-    target.addLegalOp<bufferization::ToMemrefOp, bufferization::ToTensorOp,
-                      complex::ConstantOp, complex::NotEqualOp, linalg::FillOp,
+    target.addLegalOp<complex::ConstantOp, complex::NotEqualOp, linalg::FillOp,
                       linalg::YieldOp, tensor::ExtractOp>();
     target.addLegalDialect<
         arith::ArithmeticDialect, bufferization::BufferizationDialect,
         LLVM::LLVMDialect, memref::MemRefDialect, scf::SCFDialect>();
     // Translate strategy flags to strategy options.
     SparseTensorConversionOptions options(
         sparseToSparseConversionStrategy(sparseToSparse));
     // Populate with rules and apply rewriting rules.
struct SparseTensorCodegenPass
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeToBufferConverter converter;
     ConversionTarget target(*ctx);
     // Almost everything in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
     target.addLegalOp<StorageGetOp, StorageSetOp>();
-    // All dynamic rules below accept new function, call, return.
+    // All dynamic rules below accept new function, call, return, and various
+    // tensor and bufferization operations as legal output of the rewriting
+    // provided that all sparse tensor types have been fully rewritten.
     target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
       return converter.isSignatureLegal(op.getFunctionType());
     });
     target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
       return converter.isSignatureLegal(op.getCalleeType());
     });
     target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
       return converter.isLegal(op.getOperandTypes());
     });
+    target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(
+        [&](bufferization::DeallocTensorOp op) {
+          return converter.isLegal(op.getTensor().getType());
+        });
     // Legal dialects may occur in generated code.
     target.addLegalDialect<arith::ArithmeticDialect,
                            bufferization::BufferizationDialect,
                            memref::MemRefDialect, scf::SCFDialect>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     populateCallOpTypeConversionPattern(patterns, converter);

mlir/test/Dialect/SparseTensor/codegen.mlir

 // CHECK-LABEL: func @sparse_values_dcsr(
 //  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)
 //       CHECK: %[[F:.*]] = sparse_tensor.storage_get %[[A]][5] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xf64>
 //       CHECK: return %[[F]] : memref<?xf64>
 func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
   return %0 : memref<?xf64>
 }
+
+// CHECK-LABEL: func @sparse_dealloc_csr(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)
+//       CHECK: %[[F0:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<2xindex>
+//       CHECK: memref.dealloc %[[F0]] : memref<2xindex>
+//       CHECK: %[[F1:.*]] = sparse_tensor.storage_get %[[A]][1] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi32>
+//       CHECK: memref.dealloc %[[F1]] : memref<?xi32>
+//       CHECK: %[[F2:.*]] = sparse_tensor.storage_get %[[A]][2] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi64>
+//       CHECK: memref.dealloc %[[F2]] : memref<?xi64>
+//       CHECK: %[[F3:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xf64>
+//       CHECK: memref.dealloc %[[F3]] : memref<?xf64>
+//       CHECK: return
+func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {
+  bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>
+  return
+}