[Mlir-commits] [mlir] bc7eb10 - [mlir][sparse] Add push_back op to support code generation.

[llvmlistbot at llvm.org]

Author: bixia1
Date: 2022-09-18T21:22:41-07:00
New Revision: bc7eb10cf30310a3703c9ead7dc31a3cbc141fc8

URL: https://github.com/llvm/llvm-project/commit/bc7eb10cf30310a3703c9ead7dc31a3cbc141fc8
DIFF: https://github.com/llvm/llvm-project/commit/bc7eb10cf30310a3703c9ead7dc31a3cbc141fc8.diff

LOG: [mlir][sparse] Add push_back op to support code generation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134062

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
    mlir/test/Dialect/SparseTensor/roundtrip.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index ed1943f20de2..302b56a22b15 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -232,6 +232,35 @@ def SparseTensor_InsertOp : SparseTensor_Op<"insert", []>,
                        " type($tensor) `,` type($indices) `,` type($value)";
 }
 
+def SparseTensor_PushBackOp : SparseTensor_Op<"push_back", []>,
+    Arguments<(ins StridedMemRefRankOf<[Index], [1]>:$bufferSizes,
+               StridedMemRefRankOf<[AnyType], [1]>:$inBuffer,
+               AnyType:$value, IndexAttr:$idx)>,
+    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 operation returns an SSA value for the memref. 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 -> memref<?xf64>
+    ```
+  }];
+  let assemblyFormat = "$bufferSizes `,` $inBuffer `,` $value"
+                       " attr-dict `:` type($bufferSizes) `,`"
+                       " type($inBuffer) `,` type($value) `to`"
+                       " type($outBuffer)";
+}
+
 def SparseTensor_ExpandOp : SparseTensor_Op<"expand", []>,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$values,

diff  --git a/mlir/test/Dialect/SparseTensor/roundtrip.mlir b/mlir/test/Dialect/SparseTensor/roundtrip.mlir
index 5c22ffb0cc69..7d32300c6183 100644
--- a/mlir/test/Dialect/SparseTensor/roundtrip.mlir
+++ b/mlir/test/Dialect/SparseTensor/roundtrip.mlir
@@ -132,6 +132,19 @@ func.func @sparse_insert(%arg0: tensor<128xf64, #SparseVector>, %arg1: memref<?x
 
 // -----
 
+// CHECK-LABEL: func @sparse_push_back(
+//  CHECK-SAME: %[[A:.*]]: memref<?xindex>,
+//  CHECK-SAME: %[[B:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[C:.*]]: f64) -> memref<?xf64> {
+//       CHECK: %[[D:.*]] = sparse_tensor.push_back %[[A]], %[[B]], %[[C]] {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64 to memref<?xf64>
+//       CHECK: return %[[D]]
+func.func @sparse_push_back(%arg0: memref<?xindex>, %arg1: memref<?xf64>, %arg2: f64) -> memref<?xf64> {
+  %0 = sparse_tensor.push_back %arg0, %arg1, %arg2 {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64 to memref<?xf64>
+  return %0 : memref<?xf64>
+}
+
+// -----
+
 #SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 
 // CHECK-LABEL: func @sparse_expansion(