[mlir][sparse] Extend sparse_tensor.push_back to allow push_back a value n times.

[bixia]

Fix format.

[Peiming]

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

275

Why can't we support n==1 though, it falls back to push_back, but should still be legitimate, right?

[aartbik]

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

275

Wait, is the n a compile-time constant?

That won't work? We need this as a regular SSA value.
The number of repeats will often depend on runtime computations?

[aartbik]

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

285

Perhaps also say something here that although this is semantically equivalent to calling push_back n times, it gives backend optimizations more chances to deal with memory reallocation and filling the memory with the same value.

[bixia]

Make n a runtime value.

[bixia]

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

275

There must be some miscommunication, as in the discuss the decision was to use attribute. Change to use input value.

[bixia]

Add description about the performance advantage.

[aartbik]

Neat! Thanks!

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

274

layout?

275

No worries, thanks for changing. The indent was always to have runtime values, since we want to mimic e.g.

values.insert(values.end(), count, 0);

from the runtime library in actual code.

324

Interesting, when something is optional in the assembly, there is no builder without?

[bixia]

Fix format.

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

 
 def SparseTensor_PushBackOp : SparseTensor_Op<"push_back",
     [TypesMatchWith<"value type matches element type of inBuffer",
                     "inBuffer", "value",
                     "$_self.cast<ShapedType>().getElementType()">,
      AllTypesMatch<["inBuffer", "outBuffer"]>]>,
     Arguments<(ins StridedMemRefRankOf<[Index], [1]>:$bufferSizes,
                StridedMemRefRankOf<[AnyType], [1]>:$inBuffer,
-               AnyType:$value, IndexAttr:$idx, UnitAttr:$inbounds)>,
+               AnyType:$value, IndexAttr:$idx, Optional<Index>:$n,
+               UnitAttr:$inbounds)>,
     Results<(outs StridedMemRefRankOf<[AnyType], [1]>:$outBuffer)>  {
   string summary = "Pushes a value to the back of a given buffer";
   string description = [{
     Push `value` to the end of the given sparse tensor storage buffer
     `inBuffer` and update the size of the buffer in `bufferSizes[idx]`. The
     capacity of the buffer is recorded in the memref type of `inBuffer `. If the
     current buffer is full, then `inBuffer.realloc` is called before pushing the
     data to the buffer. This is similar to std::vector push_back.
 
+    The optional input `n` specifies the number of times to repeately push
+    the value to the back of the tensor. When `n` is a compile-time constant,
+    its value can't be less than 1. If `n` is a runtime value that is less
+    than 1, the behavior is undefined. Although using input `n` is semantically
+    equivalent to calling push_back n times, it gives compiler more chances to
+    to optimize the memory reallocation and the filling of the memory with the
+    same value.
+
     The `inbounds` attribute tells the compiler that the insertion won't go
     beyond the current storage buffer. This allows the compiler to not generate
     the code for capacity check and reallocation. The typical usage will be for
     "dynamic" sparse tensors for which a capacity can be set beforehand.
 
     Note that this operation is "impure" in the sense that even though
     the result is modeled through an SSA value, referencing the memref
     through the old SSA value after this operation is undefined behavior.
 
     Example:
 
     ```mlir
     %r = sparse_tensor.push_back %bufferSizes, %buffer, %val
       {idx = 0 : index} : memref<?xindex>, memref<?xf64>, f64
     ```
 
     ```mlir
     %r = sparse_tensor.push_back inbounds %bufferSizes, %buffer, %val
       {idx = 0 : index} : memref<?xindex>, memref<?xf64>, f64
     ```
+
+    ```mlir
+    %r = sparse_tensor.push_back inbounds %bufferSizes, %buffer, %val, %n
+      {idx = 0 : index} : memref<?xindex>, memref<?xf64>, f64
+    ```
   }];
   let assemblyFormat = "(`inbounds` $inbounds^)? $bufferSizes `,` $inBuffer"
-                       " `,` $value attr-dict `:` type($bufferSizes) `,`"
-                       " type($inBuffer) `,` type($value)";
+                       " `,` $value (`,` $n^ )?  attr-dict `:`"
+                       " type($bufferSizes) `,` type($inBuffer) `,`"
+                       " type($value)  (`,` type($n)^ )?";
+
+  let builders = [
+    //Build an op without input `n`.
+    OpBuilder<(ins "Type":$outBuffer, "Value":$bufferSizes, "Value":$inBuffer,
+                   "Value":$value, "APInt":$idx)>
+  ];
+
+  let hasVerifier = 1;
 }
 
 def SparseTensor_ExpandOp : SparseTensor_Op<"expand", []>,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$values,
                   StridedMemRefRankOf<[I1],[1]>:$filled,
                   StridedMemRefRankOf<[Index],[1]>:$added,
                   Index:$count)> {

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

 
 LogicalResult InsertOp::verify() {
   RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
   if (ttp.getRank() != static_cast<int64_t>(getIndices().size()))
     return emitOpError("incorrect number of indices");
   return success();
 }
 
+void PushBackOp::build(OpBuilder &builder, OperationState &result,
+                       Type outBuffer, Value bufferSizes, Value inBuffer,
+                       Value value, APInt idx) {
+  build(builder, result, outBuffer, bufferSizes, inBuffer, value, idx, Value());
+}
+
+LogicalResult PushBackOp::verify() {
+  Value n = getN();
+  if (n) {
+    auto nValue = dyn_cast_or_null<arith::ConstantIndexOp>(n.getDefiningOp());
+    if (nValue && nValue.value() < 1)
+      return emitOpError("n must be not less than 1");
+  }
+  return success();
+}
+
 LogicalResult CompressOp::verify() {
   RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
   if (ttp.getRank() != 1 + static_cast<int64_t>(getIndices().size()))
     return emitOpError("incorrect number of indices");
   return success();
 }
 
 void ForeachOp::build(

mlir/test/Dialect/SparseTensor/invalid.mlir

 func.func @sparse_push_back(%arg0: memref<?xindex>, %arg1: memref<?xf64>, %arg2: f32) -> memref<?xf64> {
   // expected-error@+1 {{'sparse_tensor.push_back' op failed to verify that value type matches element type of inBuffer}}
   %0 = sparse_tensor.push_back %arg0, %arg1, %arg2 {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f32
   return %0 : memref<?xf64>
 }
 
 // -----
 
+func.func @sparse_push_back_n(%arg0: memref<?xindex>, %arg1: memref<?xf32>, %arg2: f32) -> memref<?xf32> {
+  %c0 = arith.constant 0: index
+  // expected-error@+1 {{'sparse_tensor.push_back' op n must be not less than 1}}
+  %0 = sparse_tensor.push_back %arg0, %arg1, %arg2, %c0 {idx = 2 : index} : memref<?xindex>, memref<?xf32>, f32, index
+  return %0 : memref<?xf32>
+}
+
+// -----
+
 func.func @sparse_unannotated_expansion(%arg0: tensor<128xf64>) {
   // expected-error@+1 {{'sparse_tensor.expand' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
   %values, %filled, %added, %count = sparse_tensor.expand %arg0
     : tensor<128xf64> to memref<?xf64>, memref<?xi1>, memref<?xindex>
   return
 }
 
 // -----

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: return %[[D]]
 func.func @sparse_push_back_inbound(%arg0: memref<?xindex>, %arg1: memref<?xf64>, %arg2: f64) -> memref<?xf64> {
   %0 = sparse_tensor.push_back inbounds %arg0, %arg1, %arg2 {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64
   return %0 : memref<?xf64>
 }
 
 // -----
 
+// CHECK-LABEL: func @sparse_push_back_n(
+//  CHECK-SAME: %[[A:.*]]: memref<?xindex>,
+//  CHECK-SAME: %[[B:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[C:.*]]: f64,
+//  CHECK-SAME: %[[D:.*]]: index) -> memref<?xf64> {
+//       CHECK: %[[E:.*]] = sparse_tensor.push_back %[[A]], %[[B]], %[[C]], %[[D]] {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64, index
+//       CHECK: return %[[E]]
+func.func @sparse_push_back_n(%arg0: memref<?xindex>, %arg1: memref<?xf64>, %arg2: f64, %arg3: index) -> memref<?xf64> {
+  %0 = sparse_tensor.push_back %arg0, %arg1, %arg2, %arg3 {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64, index
+  return %0 : memref<?xf64>
+}
+
+// -----
+
 #SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 
 // CHECK-LABEL: func @sparse_expansion(
 //  CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>)
 //       CHECK: %{{.*}}, %{{.*}}, %{{.*}}, %[[T:.*]] = sparse_tensor.expand %[[A]]
 //       CHECK: return %[[T]] : index
 func.func @sparse_expansion(%tensor: tensor<8x8xf64, #SparseMatrix>) -> index {
   %values, %filled, %added, %count = sparse_tensor.expand %tensor