[mlir][bufferization][sparse] put restriction on sparse tensor allocation

[aartbik]

Summary

Putting some direct use restrictions on tensor allocations in the
sparse case enables the use of simplifying assumptions in the
bufferization analysis.

[springerm]

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

260–261

I think you mentioned that it is also not allowed to pass an AllocTensorOp result into a ConvertOp.

[aartbik]

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

260–261

You are right that there are a few other "local" cases we should verify, but I need to think about this a bit more (and most of these are "caught" by failure to rewrite). But for now, this at least avoids passing those weird uninitialized sparse tensors across function boundaries ;-)

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let description = [{
     Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not
     alias with any other buffer, so it can be used to resolve read-after-write
     conflicts that would have been introduced by the in-place bufferization of
     another op.
 
     Both dense and sparse tensor types are supported. The result of a
     `bufferization.alloc_tensor` is a tensor value that can be used like any
     other tensor value. In practice, it is often used as the "out" operand of
-    another op. E.g.:
+    another op. Sparse tensor allocations should always be used in a local
+    construction operation and never escape the function boundary directly.
+
+    Example:
 
     ```mlir
     %c = bufferization.alloc_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
     %0 = linalg.matmul
       ins(%a, %b: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32, #SparseMatrix>)
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
+    return %0 : tensor<?x?xf32, #SparseMatrix>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
                        Optional<AnyTensor>:$copy);
 
   let results = (outs AnyTensor:$result);
 

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.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/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 //===----------------------------------------------------------------------===//
 // Helper functions
LogicalResult AllocTensorOp::verify() {
   if (getCopy() && !getDynamicSizes().empty())
     return emitError("dynamic sizes not needed when copying a tensor");
   if (!getCopy() && getType().getNumDynamicDims() !=
                         static_cast<int64_t>(getDynamicSizes().size()))
     return emitError("expected ")
            << getType().getNumDynamicDims() << " dynamic sizes";
   if (getCopy() && getCopy().getType() != getType())
     return emitError("expected that `copy` and return type match");
+
+  // For sparse tensor allocation, we require that none of its
+  // uses escapes the function boundary directly.
+  if (sparse_tensor::getSparseTensorEncoding(getType())) {
+    for (auto &use : getOperation()->getUses())
+      if (isa<func::ReturnOp, func::CallOp, func::CallIndirectOp>(
+              use.getOwner()))
+        return emitError("sparse tensor allocation should not escape function");
+  }
+
   return success();
 }
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes) {
   build(builder, result, type, dynamicSizes, /*copy=*/Value());
 }
 

mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRBufferizationDialect
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRDialect
   MLIRFuncDialect
   MLIRIR
+  MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRMemRefDialect
   )

mlir/test/Dialect/Bufferization/invalid.mlir

 }
 
 // -----
 
 func.func @escape_attr_non_bufferizable(%m0: memref<?xf32>) {
   // expected-error @+1{{'bufferization.escape' only valid on bufferizable ops}}
   %0 = memref.cast %m0 {bufferization.escape = [true]} : memref<?xf32> to memref<10xf32>
   return
-}
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
+
+func.func @sparse_alloc_direct_return() -> tensor<20x40xf32, #DCSR> {
+  // expected-error @+1{{sparse tensor allocation should not escape function}}
+  %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
+  return %0 : tensor<20x40xf32, #DCSR>
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
+
+func.func private @foo(tensor<20x40xf32, #DCSR>) -> ()
+
+func.func @sparse_alloc_call() {
+  // expected-error @+1{{sparse tensor allocation should not escape function}}
+  %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
+  call @foo(%0) : (tensor<20x40xf32, #DCSR>) -> ()
+  return
+}
+

mlir/test/Dialect/SparseTensor/conversion.mlir

 //   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I]], %[[Q]][%[[C0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[J]], %[[Q]][%[[C1]]] : memref<2xindex>
 //       CHECK: %[[NP:.*]] = llvm.mlir.null : !llvm.ptr<i8>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[Empty]], %[[NP]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_init(%arg0: index, %arg1: index) -> tensor<?x?xf64, #SparseMatrix> {
   %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #SparseMatrix>
-  return %0 : tensor<?x?xf64, #SparseMatrix>
+  %1 = sparse_tensor.load %0 : tensor<?x?xf64, #SparseMatrix>
+  return %1 : tensor<?x?xf64, #SparseMatrix>
 }
 
 // CHECK-LABEL: func @sparse_release(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: call @delSparseTensor(%[[A]]) : (!llvm.ptr<i8>) -> ()
 //       CHECK: return
 func.func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) {
   sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector>

 
 // CHECK-LABEL: func @sparse_and_dense_init(
 //       CHECK: %[[S:.*]] = call @newSparseTensor
 //       CHECK: %[[D:.*]] = bufferization.alloc_tensor
 //       CHECK: return %[[S]], %[[D]] : !llvm.ptr<i8>, tensor<?x?xf64>
 func.func @sparse_and_dense_init(%arg0: index, %arg1: index)
            -> (tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>) {
   %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #SparseMatrix>
-  %1 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64>
-  return %0, %1 : tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>
+  %1 = sparse_tensor.load %0 : tensor<?x?xf64, #SparseMatrix>
+  %2 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64>
+  return %1, %2 : tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>
 }

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

deps = [
         ":BufferizableOpInterfaceIncGen",
         ":BufferizationBaseIncGen",
         ":BufferizationOpsIncGen",
         ":CopyOpInterface",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefDialect",
+        ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         "//llvm:Support",
     ],
 )
 
 gentbl_cc_library(
     name = "BufferizationPassIncGen",