diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td @@ -85,6 +85,25 @@ 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)> { diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp --- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp +++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp @@ -235,6 +235,12 @@ 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()))) diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp @@ -99,8 +99,8 @@ /// 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(); @@ -471,6 +471,23 @@ } }; +/// Sparse conversion rule for the release operator. +class SparseTensorReleaseConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(ReleaseOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + StringRef name = "delSparseTensor"; + TypeRange none; + rewriter.create(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 { @@ -483,7 +500,7 @@ Type eltType = resType.cast().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)) @@ -514,7 +531,7 @@ Type eltType = resType.cast().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)) @@ -609,7 +626,8 @@ RewritePatternSet &patterns) { patterns.add( - typeConverter, patterns.getContext()); + SparseTensorReleaseConverter, SparseTensorToPointersConverter, + SparseTensorToIndicesConverter, SparseTensorToValuesConverter, + SparseTensorToTensorConverter>(typeConverter, + patterns.getContext()); } diff --git a/mlir/test/Dialect/SparseTensor/conversion.mlir b/mlir/test/Dialect/SparseTensor/conversion.mlir --- a/mlir/test/Dialect/SparseTensor/conversion.mlir +++ b/mlir/test/Dialect/SparseTensor/conversion.mlir @@ -112,6 +112,15 @@ return %0 : tensor } +// CHECK-LABEL: func @sparse_release( +// CHECK-SAME: %[[A:.*]]: !llvm.ptr) +// CHECK: call @delSparseTensor(%[[A]]) : (!llvm.ptr) -> () +// 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) -> !llvm.ptr // CHECK: return %[[A]] : !llvm.ptr diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir --- a/mlir/test/Dialect/SparseTensor/invalid.mlir +++ b/mlir/test/Dialect/SparseTensor/invalid.mlir @@ -8,6 +8,14 @@ // ----- +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 { %c = constant 0 : index // expected-error@+1 {{expected a sparse tensor to get pointers}} diff --git a/mlir/test/Dialect/SparseTensor/roundtrip.mlir b/mlir/test/Dialect/SparseTensor/roundtrip.mlir --- a/mlir/test/Dialect/SparseTensor/roundtrip.mlir +++ b/mlir/test/Dialect/SparseTensor/roundtrip.mlir @@ -15,6 +15,19 @@ #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, #{{.*}}> diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir @@ -95,6 +95,10 @@ %v = vector.load %m[%c0] : memref, vector<25xf64> vector.print %v : vector<25xf64> + // Release the resources. + sparse_tensor.release %a : tensor + sparse_tensor.release %x : tensor + return } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg b/mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg @@ -3,6 +3,3 @@ # No JIT on win32. if sys.platform == 'win32': config.unsupported = True - -# http://llvm.org/pr52046 -config.environment['ASAN_OPTIONS'] = 'detect_leaks=0' diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir @@ -273,6 +273,22 @@ %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 } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse-constant_to_sparse_tensor.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_constant_to_sparse_tensor.mlir rename from mlir/test/Integration/Dialect/SparseTensor/CPU/sparse-constant_to_sparse_tensor.mlir rename to mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_constant_to_sparse_tensor.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse-constant_to_sparse_tensor.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_constant_to_sparse_tensor.mlir @@ -45,6 +45,9 @@ %vr = vector.transfer_read %v[%c0], %d0: memref, vector<8xf64> vector.print %vr : vector<8xf64> + // Release the resources. + sparse_tensor.release %ts : tensor<10x8xf64, #Tensor1> + return } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir @@ -245,7 +245,17 @@ call @dumpf64(%v2) : (memref) -> () call @dumpf64(%v3) : (memref) -> () + // 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 } } - diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir @@ -84,6 +84,10 @@ : 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 } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir @@ -41,7 +41,8 @@ // 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>) { @@ -99,6 +100,7 @@ // Release the resources. memref.dealloc %xdata : memref<7x3xf64> + sparse_tensor.release %a : tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor> return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir @@ -54,7 +54,8 @@ // func @kernel_matvec(%arga: tensor, %argb: tensor, - %argx: tensor) -> tensor { + %argx: tensor {linalg.inplaceable = true}) + -> tensor { %0 = linalg.generic #matvec ins(%arga, %argb: tensor, tensor) outs(%argx: tensor) { @@ -111,6 +112,7 @@ // Release the resources. memref.dealloc %bdata : memref memref.dealloc %xdata : memref + sparse_tensor.release %a : tensor return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir @@ -40,7 +40,8 @@ func @kernel_mttkrp(%argb: tensor, %argc: tensor, %argd: tensor, - %arga: tensor) -> tensor { + %arga: tensor {linalg.inplaceable = true}) + -> tensor { %0 = linalg.generic #mttkrp ins(%argb, %argc, %argd: tensor, tensor, tensor) @@ -126,6 +127,7 @@ memref.dealloc %adata : memref memref.dealloc %cdata : memref memref.dealloc %ddata : memref + sparse_tensor.release %b : tensor return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir @@ -72,6 +72,9 @@ %v = vector.transfer_read %m[%c0], %d0: memref, vector<9xf64> vector.print %v : vector<9xf64> + // Release the resources. + sparse_tensor.release %x : tensor + return } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir @@ -71,6 +71,10 @@ : 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 } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir @@ -125,16 +125,14 @@ return %0 : tensor } - func @dump_i32(%arg0 : tensor) { - %m = memref.buffer_cast %arg0 : memref - %v = memref.load %m[] : memref + func @dump_i32(%arg0 : memref) { + %v = memref.load %arg0[] : memref vector.print %v : i32 return } - func @dump_f32(%arg0 : tensor) { - %m = memref.buffer_cast %arg0 : memref - %v = memref.load %m[] : memref + func @dump_f32(%arg0 : memref) { + %v = memref.load %arg0[] : memref vector.print %v : f32 return } @@ -203,13 +201,33 @@ // CHECK: 15 // CHECK: 10 // - call @dump_i32(%0) : (tensor) -> () - call @dump_f32(%1) : (tensor) -> () - call @dump_i32(%2) : (tensor) -> () - call @dump_f32(%3) : (tensor) -> () - call @dump_i32(%4) : (tensor) -> () - call @dump_i32(%5) : (tensor) -> () - call @dump_i32(%6) : (tensor) -> () + %m0 = memref.buffer_cast %0 : memref + call @dump_i32(%m0) : (memref) -> () + %m1 = memref.buffer_cast %1 : memref + call @dump_f32(%m1) : (memref) -> () + %m2 = memref.buffer_cast %2 : memref + call @dump_i32(%m2) : (memref) -> () + %m3 = memref.buffer_cast %3 : memref + call @dump_f32(%m3) : (memref) -> () + %m4 = memref.buffer_cast %4 : memref + call @dump_i32(%m4) : (memref) -> () + %m5 = memref.buffer_cast %5 : memref + call @dump_i32(%m5) : (memref) -> () + %m6 = memref.buffer_cast %6 : memref + call @dump_i32(%m6) : (memref) -> () + + // 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 + memref.dealloc %m1 : memref + memref.dealloc %m2 : memref + memref.dealloc %m3 : memref + memref.dealloc %m4 : memref + memref.dealloc %m5 : memref + memref.dealloc %m6 : memref return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir @@ -128,6 +128,7 @@ memref.dealloc %adata : memref memref.dealloc %bdata : memref memref.dealloc %xdata : memref + sparse_tensor.release %s : tensor return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir old mode 100644 new mode 100755 --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir @@ -28,6 +28,10 @@ // 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" ] }> @@ -163,6 +167,11 @@ 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 } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir @@ -74,6 +74,9 @@ %v = vector.transfer_read %m[%c0], %f0: memref, vector<16xf32> vector.print %v : vector<16xf32> + // Release the resources. + sparse_tensor.release %1 : tensor<8x8xf32, #CSR> + return } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir @@ -38,7 +38,7 @@ // func @kernel_spmm(%arga: tensor, %argb: tensor, - %argx: tensor) -> tensor { + %argx: tensor {linalg.inplaceable = true}) -> tensor { %0 = linalg.generic #spmm ins(%arga, %argb: tensor, tensor) outs(%argx: tensor) { @@ -101,6 +101,7 @@ // Release the resources. memref.dealloc %bdata : memref memref.dealloc %xdata : memref + sparse_tensor.release %a : tensor return } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_storage.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_storage.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_storage.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_storage.mlir @@ -250,6 +250,15 @@ %50 = vector.transfer_read %49[%c0], %d0: memref, 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 } } diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir @@ -35,7 +35,7 @@ // A kernel that sum-reduces a matrix to a single scalar. // func @kernel_sum_reduce(%arga: tensor, - %argx: tensor) -> tensor { + %argx: tensor {linalg.inplaceable = true}) -> tensor { %0 = linalg.generic #trait_sum_reduce ins(%arga: tensor) outs(%argx: tensor) { @@ -79,6 +79,7 @@ // Release the resources. memref.dealloc %xdata : memref + sparse_tensor.release %a : tensor return }