[mlir][sparse] Add optional size_hint operand to bufferization.alloc_tensor.

[aartbik]

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

57

not defined -> undefined

[bixia]

Fix the document.

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

let description = [{
     conflicts that would have been introduced by the in-place bufferization of
     another op.
 
     The optional `memory_space` attribute specifies the memory space when
     bufferizing this op. The memory space is inferred from `copy` if specified.
     If neither `copy` nor `memory_space` is specified, the default memory space
     is used during bufferization.
 
+    The optional `size_hint` operand specifies the number of non-zero elements
+    for sparse tensors. The value of `size_hint` should be not less than 1 and
+    not larger than the linear size of the corresponding dense tensor type. If
+    this requirement is not met, the behavior of the operator is undefined.
+
     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. 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>
     ```
+
+    ```mlir
+    %c = bufferization.alloc_tensor(%d1, %d2) size_hint = %noe
+      : tensor<?x?xf32, #SparseMatrix>
+    ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
                        Optional<AnyTensor>:$copy,
+                       Optional<Index>:$size_hint,
                        OptionalAttr<UI64Attr>:$memory_space);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
let extraClassDeclaration = [{
     // and return the shape.
     int64_t getStaticSize(unsigned idx) {
       assert(!isDynamicDim(idx) && "expected static size");
       return getType().getShape()[idx];
     }
   }];
 
   let builders = [
-    // Build an op without `copy` or `memory_space`.
+    // Build an op without `copy` or `memory_space` or `size_hint`.
     OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes)>,
 
-    // Build an op without `memory_space`.
+    // Build an op without `memory_space` or `size_hint`.
     OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
                    "Value":$copy)>,
+
+    // Build an op without `size_hint`.
+    OpBuilder<(ins "TensorType":$type, "ValueRange":$dynamicSizes,
+                   "Value":$copy, "IntegerAttr":$memory_space)>,
   ];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

LogicalResult AllocTensorOp::verify() {
   }
 
   return success();
 }
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes) {
   build(builder, result, type, dynamicSizes, /*copy=*/Value(),
+        /*size_hint=*/Value(),
         /*memory_space=*/IntegerAttr());
 }
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes,
                           Value copy) {
-  build(builder, result, type, dynamicSizes, copy,
+  build(builder, result, type, dynamicSizes, copy, /*size_hint=*/Value(),
         /*memory_space=*/IntegerAttr());
 }
 
+void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
+                          TensorType type, ValueRange dynamicSizes, Value copy,
+                          IntegerAttr memorySpace) {
+  build(builder, result, type, dynamicSizes, copy, /*size_hint=*/Value(),
+        memorySpace);
+}
+
 namespace {
 /// Change the type of the result of a `bufferization.alloc_tensor` by making
 /// the result type statically sized along dimension that in the original
 /// operation where defined as dynamic, but the size was defined using a
 /// `constant` op. For example:
 ///
 ///  %c5 = arith.constant 5: index
 ///  %0 = bufferization.alloc_tensor(%arg0, %c5) : tensor<?x?xf32>
if (parser.parseLParen() || parser.parseOperandList(dynamicSizesOperands) ||
       parser.parseRParen())
     return failure();
   ParseResult copyKeyword = parser.parseOptionalKeyword("copy");
   OpAsmParser::UnresolvedOperand copyOperand;
   if (copyKeyword.succeeded())
     if (parser.parseLParen() || parser.parseOperand(copyOperand) ||
         parser.parseRParen())
       return failure();
+  ParseResult sizeHintKeyword = parser.parseOptionalKeyword("size_hint");
+  OpAsmParser::UnresolvedOperand sizeHintOperand;
+  if (sizeHintKeyword.succeeded())
+    if (parser.parseEqual() || parser.parseOperand(sizeHintOperand))
+      return failure();
   if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon())
     return failure();
 
   TensorType type;
   if (parser.parseCustomTypeWithFallback(type))
     return failure();
   result.addTypes(type);
 
   Type indexType = parser.getBuilder().getIndexType();
   if (parser.resolveOperands(dynamicSizesOperands, indexType, result.operands))
     return failure();
   if (copyKeyword.succeeded())
     if (parser.resolveOperand(copyOperand, type, result.operands))
       return failure();
+  if (sizeHintKeyword.succeeded())
+    if (parser.resolveOperand(sizeHintOperand, indexType, result.operands))
+      return failure();
   result.addAttribute(AllocTensorOp::getOperandSegmentSizeAttr(),
                       parser.getBuilder().getDenseI32ArrayAttr(
                           {static_cast<int32_t>(dynamicSizesOperands.size()),
-                           static_cast<int32_t>(copyKeyword.succeeded())}));
+                           static_cast<int32_t>(copyKeyword.succeeded()),
+                           static_cast<int32_t>(sizeHintKeyword.succeeded())}));
   return success();
 }
 
 void AllocTensorOp::print(OpAsmPrinter &p) {
   p << "(" << getDynamicSizes() << ")";
   if (getCopy())
     p << " copy(" << getCopy() << ")";
+  if (getSizeHint())
+    p << " size_hint=" << getSizeHint();
   p.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{
                               AllocTensorOp::getOperandSegmentSizeAttr()});
   p << " : ";
   auto type = getResult().getType();
   if (auto validType = type.dyn_cast<::mlir::TensorType>())
     p.printStrippedAttrOrType(validType);
   else
     p << type;

mlir/test/Dialect/Bufferization/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = ["dense", "compressed"]
+}>
+
 // CHECK-LABEL: func @test_clone
 func.func @test_clone(%buf : memref<*xf32>) -> memref<*xf32> {
   %clone = bufferization.clone %buf : memref<*xf32> to memref<*xf32>
   return %clone : memref<*xf32>
 }
 
 // CHECK-LABEL: test_to_memref
 func.func @test_to_memref(%arg0: tensor<?xi64>, %arg1: tensor<*xi64>)
func.func @test_alloc_tensor_op(%t: tensor<?x5xf32>, %sz: index)
   // CHECK: bufferization.alloc_tensor() : tensor<5x6xf32>
   %2 = bufferization.alloc_tensor() : tensor<5x6xf32>
   // CHECK: bufferization.alloc_tensor(%{{.*}}, %{{.*}}) : tensor<?x?xf32>
   %3 = bufferization.alloc_tensor(%sz, %sz) : tensor<?x?xf32>
   // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<?x5xf32>
   %4 = bufferization.alloc_tensor() copy(%t) {escape = true} : tensor<?x5xf32>
   // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = false} : tensor<?x5xf32>
   %5 = bufferization.alloc_tensor() copy(%t) {escape = false} : tensor<?x5xf32>
+  %c100 = arith.constant 100 : index
+  // CHECK: bufferization.alloc_tensor() size_hint=
+  %6 = bufferization.alloc_tensor() size_hint=%c100 : tensor<100x100xf64, #CSR>
   return %1 : tensor<?x5xf32>
 }
 
 // CHECK-LABEL: func @test_dealloc_tensor_op
 func.func @test_dealloc_tensor_op(%arg0: tensor<4xi32>) {
   // CHECK: bufferization.dealloc_tensor {{.*}} : tensor<4xi32>
   bufferization.dealloc_tensor %arg0 : tensor<4xi32>
   return
 }