[mlir][sparse] add a "release" operation to sparse tensor dialect

[aartbik]

Summary

We have several ways to materialize sparse tensors (new and convert) but no explicit operation to release the underlying sparse storage scheme at runtime (other than making an explicit delSparseTensor() library call). To simplify memory management, a sparse_tensor.release operation has been introduced that lowers to the runtime library call while keeping tensors, opague pointers, and memrefs transparent in the initial IR.

*Note* There is obviously some tension between the concept of immutable tensors and memory management methods. This tension is addressed by simply stating that after the "release" call, no further memref related operations are allowed on the tensor value. We expect the design to evolve over time, however, and arrive at a more satisfactory view of tensors and buffers eventually.

Bug:
http://llvm.org/pr52046

[aartbik]

add noasan to one test

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

string description = [{
     %2 = sparse_tensor.convert %3 : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR>
     ```
 
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
   let hasFolder = 1;
 }
 
+def SparseTensor_ReleaseOp : SparseTensor_Op<"release", []>,
+    Arguments<(ins AnyTensor:$tensor)> {
+  string description = [{
+    Releases the underlying sparse storage scheme for a tensor that
+    materialized earlier through a `new` operator or a non-trivial
+    `convert` operator with an annotated tensor type as destination.
+    This operation should only be called once for any materialized tensor.
+    Also, after this operation, any subsequent `memref` querying operation
+    on the tensor returns undefined results.
+
+    Example:
+
+    ```mlir
+    sparse_tensor.release %tensor : tensor<1024x1024xf64, #CSR>
+    ```
+  }];
+  let assemblyFormat = "$tensor attr-dict `:` type($tensor)";
+}
+
 def SparseTensor_ToPointersOp : SparseTensor_Op<"pointers", [NoSideEffect]>,
     Arguments<(ins AnyTensor:$tensor, Index:$dim)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extract pointers array at given dimension from a tensor";
   let description = [{
     Returns the pointers array of the sparse storage scheme at the
     given dimension for the given sparse tensor. This is similar to the
     `memref.buffer_cast` operation in the sense that it provides a bridge

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 }
 
 OpFoldResult ConvertOp::fold(ArrayRef<Attribute> operands) {
   if (getType() == source().getType())
     return source();
   return {};
 }
 
+static LogicalResult verify(ReleaseOp op) {
+  if (!getSparseTensorEncoding(op.tensor().getType()))
+    return op.emitError("expected a sparse tensor to release");
+  return success();
+}
+
 static LogicalResult verify(ToPointersOp op) {
   if (auto e = getSparseTensorEncoding(op.tensor().getType())) {
     if (failed(isInBounds(op.dim(), op.tensor())))
       return op.emitError("requested pointers dimension out of bounds");
     if (failed(isMatchingWidth(op.result(), e.getPointerBitWidth())))
       return op.emitError("unexpected type for pointers");
     return success();
   }

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

auto elts =
       rewriter.create<ConstantOp>(loc, DenseElementsAttr::get(tt1, values));
   return rewriter.create<tensor::CastOp>(loc, tt2, elts);
 }
 
 /// Returns a function reference (first hit also inserts into module). Sets
 /// the "_emit_c_interface" on the function declaration when requested,
 /// so that LLVM lowering generates a wrapper function that takes care
 /// of ABI complications with passing in and returning MemRefs to C functions.
-static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type resultType,
-                                 ValueRange operands,
+static FlatSymbolRefAttr getFunc(Operation *op, StringRef name,
+                                 TypeRange resultType, ValueRange operands,
                                  bool emitCInterface = false) {
   MLIRContext *context = op->getContext();
   auto module = op->getParentOfType<ModuleOp>();
   auto result = SymbolRefAttr::get(context, name);
   auto func = module.lookupSymbol<FuncOp>(result.getAttr());
   if (!func) {
     OpBuilder moduleBuilder(module.getBodyRegion());
     func = moduleBuilder.create<FuncOp>(
scf::buildLoopNest(rewriter, op.getLoc(), lo, hi, st, {},
                                        perm);
                          return {};
                        });
     rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, ptr));
     return success();
   }
 };
 
+/// Sparse conversion rule for the release operator.
+class SparseTensorReleaseConverter : public OpConversionPattern<ReleaseOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ReleaseOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    StringRef name = "delSparseTensor";
+    TypeRange none;
+    rewriter.create<CallOp>(op.getLoc(), none,
+                            getFunc(op, name, none, adaptor.getOperands()),
+                            adaptor.getOperands());
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 /// Sparse conversion rule for pointer accesses.
 class SparseTensorToPointersConverter
     : public OpConversionPattern<ToPointersOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToPointersOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isIndex())
-      name = "sparsePointers";
+      name = "sparsePointers"; // 64-bit, but its own name for unique signature
     else if (eltType.isInteger(64))
       name = "sparsePointers64";
     else if (eltType.isInteger(32))
       name = "sparsePointers32";
     else if (eltType.isInteger(16))
       name = "sparsePointers16";
     else if (eltType.isInteger(8))
       name = "sparsePointers8";
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isIndex())
-      name = "sparseIndices";
+      name = "sparseIndices"; // 64-bit, but its own name for unique signature
     else if (eltType.isInteger(64))
       name = "sparseIndices64";
     else if (eltType.isInteger(32))
       name = "sparseIndices32";
     else if (eltType.isInteger(16))
       name = "sparseIndices16";
     else if (eltType.isInteger(8))
       name = "sparseIndices8";

 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
                                                   RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
                SparseTensorNewConverter, SparseTensorConvertConverter,
-               SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
-               SparseTensorToValuesConverter, SparseTensorToTensorConverter>(
-      typeConverter, patterns.getContext());
+               SparseTensorReleaseConverter, SparseTensorToPointersConverter,
+               SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
+               SparseTensorToTensorConverter>(typeConverter,
+                                              patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/conversion.mlir

 //   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<3xi64> to tensor<?xi64>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor>
   return %0 : tensor<?x?x?xf32, #SparseTensor>
 }
 
+// CHECK-LABEL: func @sparse_release(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: call @delSparseTensor(%[[A]]) : (!llvm.ptr<i8>) -> ()
+//       CHECK: return
+func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) {
+  sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector>
+  return
+}
+
 // CHECK-LABEL: func @sparse_nop_convert(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 func @sparse_nop_convert(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector>
   return %0 : tensor<64xf32, #SparseVector>
 }
 

mlir/test/Dialect/SparseTensor/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 func @invalid_new_dense(%arg0: !llvm.ptr<i8>) -> tensor<32xf32> {
   // expected-error@+1 {{expected a sparse tensor result}}
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
+func @invalid_release_dense(%arg0: tensor<4xi32>) {
+  // expected-error@+1 {{expected a sparse tensor to release}}
+  sparse_tensor.release %arg0 : tensor<4xi32>
+  return
+}
+
+// -----
+
 func @invalid_pointers_dense(%arg0: tensor<128xf64>) -> memref<?xindex> {
   %c = constant 0 : index
   // expected-error@+1 {{expected a sparse tensor to get pointers}}
   %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
 
 // -----

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
   return %0 : tensor<128xf64, #SparseVector>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
+// CHECK-LABEL: func @sparse_release(
+// CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>
+//       CHECK: sparse_tensor.release %[[A]] : tensor<128xf64, #{{.*}}>
+//       CHECK: return
+func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) {
+  sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector>
+  return
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
 // CHECK-LABEL: func @sparse_convert_1d_to_sparse(
 // CHECK-SAME: %[[A:.*]]: tensor<64xf32>)
 //       CHECK: %[[T:.*]] = sparse_tensor.convert %[[A]] : tensor<64xf32> to tensor<64xf32, #{{.*}}>
 //       CHECK: return %[[T]] : tensor<64xf32, #{{.*}}>
 func @sparse_convert_1d_to_sparse(%arg0: tensor<64xf32>) -> tensor<64xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<64xf32> to tensor<64xf32, #SparseVector>
   return %0 : tensor<64xf32, #SparseVector>
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir

func @entry() {
     //
     // CHECK: ( 1, 0, 0, 1.4, 0, 0, 2, 0, 0, 2.5, 0, 0, 3, 0, 0, 4.1, 0, 0, 4, 0, 0, 5.2, 0, 0, 5 )
     //
     %m = sparse_tensor.values %0
       : tensor<?x?xf64, #DenseMatrix> to memref<?xf64>
     %v = vector.load %m[%c0] : memref<?xf64>, vector<25xf64>
     vector.print %v : vector<25xf64>
 
+    // Release the resources.
+    sparse_tensor.release %a : tensor<?x?xf64, #SparseMatrix>
+    sparse_tensor.release %x : tensor<?x?xf64, #DenseMatrix>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg

 import sys
 
 # No JIT on win32.
 if sys.platform == 'win32':
     config.unsupported = True
-
-# http://llvm.org/pr52046
-config.environment['ASAN_OPTIONS'] = 'detect_leaks=0'

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

func @entry() {
     //
     // CHECK: ( -1064514355, -1068289229, -1072902963, -1081291571, 0, 1066192077, 1074580685, 1079194419, 1082969293, 1134084096 )
     //
     %c9 = call @sparse_cast_f32_as_s32(%3) : (tensor<10xf32, #SV>) -> tensor<10xi32>
     %m9 = memref.buffer_cast %c9 : memref<10xi32>
     %v9 = vector.transfer_read %m9[%z], %i: memref<10xi32>, vector<10xi32>
     vector.print %v9 : vector<10xi32>
 
+    // Release the resources.
+    sparse_tensor.release %1 : tensor<10xi32, #SV>
+    sparse_tensor.release %3 : tensor<10xf32, #SV>
+    sparse_tensor.release %5 : tensor<10xf64, #SV>
+    sparse_tensor.release %7 : tensor<10xf64, #SV>
+    memref.dealloc %m0 : memref<10xf32>
+    memref.dealloc %m1 : memref<10xf32>
+    memref.dealloc %m2 : memref<10xi32>
+    memref.dealloc %m3 : memref<10xi32>
+    memref.dealloc %m4 : memref<10xf64>
+    memref.dealloc %m5 : memref<10xf32>
+    memref.dealloc %m6 : memref<10xi64>
+    memref.dealloc %m7 : memref<10xi64>
+    memref.dealloc %m8 : memref<10xi8>
+    memref.dealloc %m9 : memref<10xi32>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_constant_to_sparse_tensor.mlir

func @entry() {
     %i1r = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<8xindex>
     vector.print %i1r : vector<8xindex>
 
     // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8 )
     %v = sparse_tensor.values %ts : tensor<10x8xf64, #Tensor1> to memref<?xf64>
     %vr = vector.transfer_read %v[%c0], %d0: memref<?xf64>, vector<8xf64>
     vector.print %vr : vector<8xf64>
 
+    // Release the resources.
+    sparse_tensor.release %ts : tensor<10x8xf64, #Tensor1>
+
     return
   }
 }
 

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir

func @entry() {
     // CHECK:      ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 )
     // CHECK-NEXT: ( 1, 13, 2, 14, 3, 15, 4, 16, 5, 17, 6, 18, 7, 19, 8, 20, 9, 21, 10, 22, 11, 23, 12, 24 )
     // CHECK-NEXT: ( 1, 5, 9, 13, 17, 21, 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23, 4, 8, 12, 16, 20, 24 )
     //
     call @dumpf64(%v1) : (memref<?xf64>) -> ()
     call @dumpf64(%v2) : (memref<?xf64>) -> ()
     call @dumpf64(%v3) : (memref<?xf64>) -> ()
 
+    // Release the resources.
+    sparse_tensor.release %1 : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %2 : tensor<2x3x4xf64, #Tensor2>
+    sparse_tensor.release %3 : tensor<2x3x4xf64, #Tensor3>
+    sparse_tensor.release %b : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %c : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %d : tensor<2x3x4xf64, #Tensor2>
+    sparse_tensor.release %f : tensor<2x3x4xf64, #Tensor2>
+    sparse_tensor.release %g : tensor<2x3x4xf64, #Tensor3>
+    sparse_tensor.release %h : tensor<2x3x4xf64, #Tensor3>
+
     return
   }
 }
-

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

func @entry() {
     // CHECK-SAME: ( 0, 0, 3, 6, -3, -6 ),
     // CHECK-SAME: ( 2, -1, 3, 0, -3, 0 ) )
     //
     %m = memref.buffer_cast %0 : memref<6x6xi32>
     %v = vector.transfer_read %m[%c0, %c0], %i0
       : memref<6x6xi32>, vector<6x6xi32>
     vector.print %v : vector<6x6xi32>
 
+    // Release the resources.
+    sparse_tensor.release %sparse_filter : tensor<3x3xi32, #DCSR>
+    memref.dealloc %m : memref<6x6xi32>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

 // actual sparse code, initializes a matching sparse storage scheme
 // from file, and runs the resulting code with the JIT compiler.
 //
 module {
   //
   // A kernel that flattens a rank 8 tensor into a dense matrix.
   //
   func @kernel_flatten(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>,
-                       %argx: tensor<7x3xf64>) -> tensor<7x3xf64> {
+                       %argx: tensor<7x3xf64> {linalg.inplaceable = true})
+           -> tensor<7x3xf64> {
     %0 = linalg.generic #trait_flatten
       ins(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>)
       outs(%argx: tensor<7x3xf64>) {
       ^bb(%a: f64, %x: f64):
         %0 = addf %x, %a : f64
         linalg.yield %0 : f64
     } -> tensor<7x3xf64>
     return %0 : tensor<7x3xf64>
func @entry() {
     %r = memref.buffer_cast %0 : memref<7x3xf64>
     scf.for %i = %c0 to %c7 step %c1 {
       %v = vector.transfer_read %r[%i, %c0], %d0: memref<7x3xf64>, vector<3xf64>
       vector.print %v : vector<3xf64>
     }
 
     // Release the resources.
     memref.dealloc %xdata : memref<7x3xf64>
+    sparse_tensor.release %a : tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

 //
 module {
   //
   // A kernel that multiplies a sparse matrix A with a dense vector b
   // into a dense vector x.
   //
   func @kernel_matvec(%arga: tensor<?x?xi32, #SparseMatrix>,
                       %argb: tensor<?xi32>,
-                      %argx: tensor<?xi32>) -> tensor<?xi32> {
+                      %argx: tensor<?xi32> {linalg.inplaceable = true})
+          -> tensor<?xi32> {
     %0 = linalg.generic #matvec
       ins(%arga, %argb: tensor<?x?xi32, #SparseMatrix>, tensor<?xi32>)
       outs(%argx: tensor<?xi32>) {
       ^bb(%a: i32, %b: i32, %x: i32):
         %0 = muli %a, %b : i32
         %1 = addi %x, %0 : i32
         linalg.yield %1 : i32
     } -> tensor<?xi32>
func @entry() {
     //
     %m = memref.buffer_cast %0 : memref<?xi32>
     %v = vector.transfer_read %m[%c0], %i0: memref<?xi32>, vector<4xi32>
     vector.print %v : vector<4xi32>
 
     // Release the resources.
     memref.dealloc %bdata : memref<?xi32>
     memref.dealloc %xdata : memref<?xi32>
+    sparse_tensor.release %a : tensor<?x?xi32, #SparseMatrix>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir

 module {
   //
   // Computes Matricized Tensor Times Khatri-Rao Product (MTTKRP) kernel. See
   // http://tensor-compiler.org/docs/data_analytics/index.html.
   //
   func @kernel_mttkrp(%argb: tensor<?x?x?xf64, #SparseMatrix>,
                       %argc: tensor<?x?xf64>,
                       %argd: tensor<?x?xf64>,
-                      %arga: tensor<?x?xf64>) -> tensor<?x?xf64> {
+                      %arga: tensor<?x?xf64> {linalg.inplaceable = true})
+          -> tensor<?x?xf64> {
     %0 = linalg.generic #mttkrp
       ins(%argb, %argc, %argd:
             tensor<?x?x?xf64, #SparseMatrix>, tensor<?x?xf64>, tensor<?x?xf64>)
       outs(%arga: tensor<?x?xf64>) {
       ^bb(%b: f64, %c: f64, %d: f64, %a: f64):
         %0 = mulf %b, %c : f64
         %1 = mulf %d, %0 : f64
         %2 = addf %a, %1 : f64
func @entry() {
     %v = vector.transfer_read %m[%c0, %c0], %i0
           : memref<?x?xf64>, vector<2x5xf64>
     vector.print %v : vector<2x5xf64>
 
     // Release the resources.
     memref.dealloc %adata : memref<?x?xf64>
     memref.dealloc %cdata : memref<?x?xf64>
     memref.dealloc %ddata : memref<?x?xf64>
+    sparse_tensor.release %b : tensor<?x?x?xf64, #SparseMatrix>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir

func @entry() {
     // Print the result for verification.
     //
     // CHECK: ( 1, 1.96, 4, 6.25, 9, 16.81, 16, 27.04, 25 )
     //
     %m = sparse_tensor.values %0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
     %v = vector.transfer_read %m[%c0], %d0: memref<?xf64>, vector<9xf64>
     vector.print %v : vector<9xf64>
 
+    // Release the resources.
+    sparse_tensor.release %x : tensor<?x?xf64, #DCSR>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir

func @entry() {
     // CHECK-SAME: ( -254, 0, 256, -300, -30, -6 ),
     // CHECK-SAME: ( 1397, 0, -1408, 100, 10, 33 ) )
     //
     %m = memref.buffer_cast %0 : memref<5x6xi32>
     %v = vector.transfer_read %m[%c0, %c0], %i0
       : memref<5x6xi32>, vector<5x6xi32>
     vector.print %v : vector<5x6xi32>
 
+    // Release the resources.
+    sparse_tensor.release %sparse_input2 : tensor<3x6xi8, #DCSR>
+    memref.dealloc %m : memref<5x6xi32>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir

%0 = linalg.generic #trait_reduction
       outs(%argx: tensor<i32>) {
         ^bb(%a: i32, %x: i32):
           %0 = xor %x, %a : i32
           linalg.yield %0 : i32
     } -> tensor<i32>
     return %0 : tensor<i32>
   }
 
-  func @dump_i32(%arg0 : tensor<i32>) {
-    %m = memref.buffer_cast %arg0 : memref<i32>
-    %v = memref.load %m[] : memref<i32>
+  func @dump_i32(%arg0 : memref<i32>) {
+    %v = memref.load %arg0[] : memref<i32>
     vector.print %v : i32
     return
   }
 
-  func @dump_f32(%arg0 : tensor<f32>) {
-    %m = memref.buffer_cast %arg0 : memref<f32>
-    %v = memref.load %m[] : memref<f32>
+  func @dump_f32(%arg0 : memref<f32>) {
+    %v = memref.load %arg0[] : memref<f32>
     vector.print %v : f32
     return
   }
 
   func @entry() {
     %ri = constant dense< 7   > : tensor<i32>
     %rf = constant dense< 2.0 > : tensor<f32>
 
func @entry() {
     // CHECK: 26
     // CHECK: 27.5
     // CHECK: 3087
     // CHECK: 168
     // CHECK: 1
     // CHECK: 15
     // CHECK: 10
     //
-    call @dump_i32(%0) : (tensor<i32>) -> ()
-    call @dump_f32(%1) : (tensor<f32>) -> ()
-    call @dump_i32(%2) : (tensor<i32>) -> ()
-    call @dump_f32(%3) : (tensor<f32>) -> ()
-    call @dump_i32(%4) : (tensor<i32>) -> ()
-    call @dump_i32(%5) : (tensor<i32>) -> ()
-    call @dump_i32(%6) : (tensor<i32>) -> ()
+    %m0 = memref.buffer_cast %0 : memref<i32>
+    call @dump_i32(%m0) : (memref<i32>) -> ()
+    %m1 = memref.buffer_cast %1 : memref<f32>
+    call @dump_f32(%m1) : (memref<f32>) -> ()
+    %m2 = memref.buffer_cast %2 : memref<i32>
+    call @dump_i32(%m2) : (memref<i32>) -> ()
+    %m3 = memref.buffer_cast %3 : memref<f32>
+    call @dump_f32(%m3) : (memref<f32>) -> ()
+    %m4 = memref.buffer_cast %4 : memref<i32>
+    call @dump_i32(%m4) : (memref<i32>) -> ()
+    %m5 = memref.buffer_cast %5 : memref<i32>
+    call @dump_i32(%m5) : (memref<i32>) -> ()
+    %m6 = memref.buffer_cast %6 : memref<i32>
+    call @dump_i32(%m6) : (memref<i32>) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sparse_input_i32 : tensor<32xi32, #SV>
+    sparse_tensor.release %sparse_input_f32 : tensor<32xf32, #SV>
+    sparse_tensor.release %dense_input_i32  : tensor<32xi32, #DV>
+    sparse_tensor.release %dense_input_f32  : tensor<32xf32, #DV>
+    memref.dealloc %m0 : memref<i32>
+    memref.dealloc %m1 : memref<f32>
+    memref.dealloc %m2 : memref<i32>
+    memref.dealloc %m3 : memref<f32>
+    memref.dealloc %m4 : memref<i32>
+    memref.dealloc %m5 : memref<i32>
+    memref.dealloc %m6 : memref<i32>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir

scf.for %i = %c0 to %c5 step %c1 {
       %v = vector.transfer_read %r[%i, %c0], %d0: memref<?x?xf32>, vector<5xf32>
       vector.print %v : vector<5xf32>
     }
 
     // Release the resources.
     memref.dealloc %adata : memref<?x?xf32>
     memref.dealloc %bdata : memref<?x?xf32>
     memref.dealloc %xdata : memref<?x?xf32>
+    sparse_tensor.release %s : tensor<?x?xf32, #SparseMatrix>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir

 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
 // RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-math-to-llvm \
 // RUN:   --convert-std-to-llvm --reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// Interop between linalg/sparse leaves some issues to be revolved:
+// UNSUPPORTED: asan
+
 
 #SM = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 #trait_sampled_dense_dense = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i,j)>,  // S
     affine_map<(i,j,k) -> (i,k)>,  // A
     affine_map<(i,j,k) -> (k,j)>,  // B
func @entry() {
     %m1 = memref.buffer_cast %1 : memref<8x8xf64>
     %v0 = vector.transfer_read %m0[%c0, %c0], %d0
         : memref<8x8xf64>, vector<8x8xf64>
     %v1 = vector.transfer_read %m1[%c0, %c0], %d0
         : memref<8x8xf64>, vector<8x8xf64>
     vector.print %v0 : vector<8x8xf64>
     vector.print %v1 : vector<8x8xf64>
 
+    // Release the resources.
+    sparse_tensor.release %s : tensor<8x8xf64, #SM>
+    memref.dealloc %m0 : memref<8x8xf64>
+    memref.dealloc %m1 : memref<8x8xf64>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir

func @entry() {
     // Print the resulting compacted values for verification.
     //
     // CHECK: ( 2, 2, 2, 4, 6, 8, 2, 10, 2, 2, 12, 2, 14, 2, 2, 16 )
     //
     %m = sparse_tensor.values %2 : tensor<8x8xf32, #CSR> to memref<?xf32>
     %v = vector.transfer_read %m[%c0], %f0: memref<?xf32>, vector<16xf32>
     vector.print %v : vector<16xf32>
 
+    // Release the resources.
+    sparse_tensor.release %1 : tensor<8x8xf32, #CSR>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir

 //
 module {
   //
   // A kernel that multiplies a sparse matrix A with a dense matrix B
   // into a dense matrix X.
   //
   func @kernel_spmm(%arga: tensor<?x?xf64, #SparseMatrix>,
                     %argb: tensor<?x?xf64>,
-                    %argx: tensor<?x?xf64>) -> tensor<?x?xf64> {
+                    %argx: tensor<?x?xf64> {linalg.inplaceable = true}) -> tensor<?x?xf64> {
     %0 = linalg.generic #spmm
       ins(%arga, %argb: tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>)
       outs(%argx: tensor<?x?xf64>) {
       ^bb(%a: f64, %b: f64, %x: f64):
         %0 = mulf %a, %b : f64
         %1 = addf %x, %0 : f64
         linalg.yield %1 : f64
     } -> tensor<?x?xf64>
func @entry() {
     //
     %m = memref.buffer_cast %0 : memref<?x?xf64>
     %v = vector.transfer_read %m[%c0, %c0], %i0: memref<?x?xf64>, vector<4x4xf64>
     vector.print %v : vector<4x4xf64>
 
     // Release the resources.
     memref.dealloc %bdata : memref<?x?xf64>
     memref.dealloc %xdata : memref<?x?xf64>
+    sparse_tensor.release %a : tensor<?x?xf64, #SparseMatrix>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_storage.mlir

func @entry() {
     vector.print %46 : vector<2xindex>
     %47 = sparse_tensor.indices %y, %c0 : tensor<10x8xf64, #BlockCol> to memref<?xindex>
     %48 = vector.transfer_read %47[%c0], %c0: memref<?xindex>, vector<7xindex>
     vector.print %48 : vector<7xindex>
     %49 = sparse_tensor.values %y : tensor<10x8xf64, #BlockCol> to memref<?xf64>
     %50 = vector.transfer_read %49[%c0], %d0: memref<?xf64>, vector<70xf64>
     vector.print %50 : vector<70xf64>
 
+    // Release the resources.
+    sparse_tensor.release %0 : tensor<10x8xf64, #Dense>
+    sparse_tensor.release %1 : tensor<10x8xf64, #CSR>
+    sparse_tensor.release %2 : tensor<10x8xf64, #DCSR>
+    sparse_tensor.release %3 : tensor<10x8xf64, #CSC>
+    sparse_tensor.release %4 : tensor<10x8xf64, #DCSC>
+    sparse_tensor.release %x : tensor<10x8xf64, #BlockRow>
+    sparse_tensor.release %y : tensor<10x8xf64, #BlockCol>
+
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

 // actual sparse code, initializes a matching sparse storage scheme
 // from file, and runs the resulting code with the JIT compiler.
 //
 module {
   //
   // A kernel that sum-reduces a matrix to a single scalar.
   //
   func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix>,
-                          %argx: tensor<f64>) -> tensor<f64> {
+                          %argx: tensor<f64> {linalg.inplaceable = true}) -> tensor<f64> {
     %0 = linalg.generic #trait_sum_reduce
       ins(%arga: tensor<?x?xf64, #SparseMatrix>)
       outs(%argx: tensor<f64>) {
       ^bb(%a: f64, %x: f64):
         %0 = addf %x, %a : f64
         linalg.yield %0 : f64
     } -> tensor<f64>
     return %0 : tensor<f64>
func @entry() {
     // CHECK: 28.2
     //
     %m = memref.buffer_cast %0 : memref<f64>
     %v = memref.load %m[] : memref<f64>
     vector.print %v : f64
 
     // Release the resources.
     memref.dealloc %xdata : memref<f64>
+    sparse_tensor.release %a : tensor<?x?xf64, #SparseMatrix>
 
     return
   }
 }