[mlir][sparse] Allow the push_back operator to skip capacity check and reallocation.

[bixia]

Summary

Add UnitAttr inbounds for this purpose.

[aartbik]

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

255

Make a comment that the typical usage will be for "dynamic" sparse tensors for which a capacity can be set beforehand

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

354

is present

[bixia]

Fix comments to address review feedback.

[bixia]

arc land can't determine revision.

[bixia]

Rebase

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

def SparseTensor_InsertOp : SparseTensor_Op<"insert",
   }];
   let assemblyFormat = "$value `into` $tensor `[` $indices `]` attr-dict`:` type($tensor)";
   let hasVerifier = 1;
 }
 
 def SparseTensor_PushBackOp : SparseTensor_Op<"push_back", []>,
     Arguments<(ins StridedMemRefRankOf<[Index], [1]>:$bufferSizes,
                StridedMemRefRankOf<[AnyType], [1]>:$inBuffer,
-               AnyType:$value, IndexAttr:$idx)>,
+               AnyType:$value, IndexAttr:$idx, 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 `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.
+
     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>
     ```
+
+    ```mlir
+    %r = sparse_tensor.push_back inbounds %bufferSizes, %buffer, %val
+      {idx = 0 : index} : memref<?xindex>, memref<?xf64>, f64 -> memref<?xf64>
+    ```
   }];
-  let assemblyFormat = "$bufferSizes `,` $inBuffer `,` $value"
-                       " attr-dict `:` type($bufferSizes) `,`"
+  let assemblyFormat = "(`inbounds` $inbounds^)? $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,
                   StridedMemRefRankOf<[I1],[1]>:$filled,

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

LogicalResult matchAndRewrite(PushBackOp op,
                                 PatternRewriter &rewriter) const override {
     // Rewrite push_back(buffer, value) to:
     // if (size(buffer) >= capacity(buffer))
     //    new_capacity = capacity(buffer)*2
     //    new_buffer = realloc(buffer, new_capacity)
     // buffer = new_buffer
     // store(buffer, value)
     // size(buffer)++
+    //
+    // The capacity check is skipped when the attribute inbounds is presented.
     Location loc = op->getLoc();
     Value c0 = constantIndex(rewriter, loc, 0);
     Value buffer = op.getInBuffer();
     Value capacity = rewriter.create<memref::DimOp>(loc, buffer, c0);
     Value idx = constantIndex(rewriter, loc, op.getIdx().getZExtValue());
     Value bufferSizes = op.getBufferSizes();
     Value size = rewriter.create<memref::LoadOp>(loc, bufferSizes, idx);
-    Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge,
-                                                size, capacity);
     Value value = op.getValue();
+
+    if (!op.getInbounds()) {
+      Value cond = rewriter.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::uge, size, capacity);
+
       auto bufferType =
           MemRefType::get({ShapedType::kDynamicSize}, value.getType());
       scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, bufferType, cond,
                                                   /*else=*/true);
       // True branch.
       rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
       Value c2 = constantIndex(rewriter, loc, 2);
       capacity = rewriter.create<arith::MulIOp>(loc, capacity, c2);
       Value newBuffer =
           rewriter.create<memref::ReallocOp>(loc, bufferType, buffer, capacity);
       rewriter.create<scf::YieldOp>(loc, newBuffer);
 
       // False branch.
       rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
       rewriter.create<scf::YieldOp>(loc, buffer);
 
-    // Add the value to the end of the buffer.
+      // Prepare for adding the value to the end of the buffer.
       rewriter.setInsertionPointAfter(ifOp);
       buffer = ifOp.getResult(0);
+    }
+
+    // Add the value to the end of the buffer.
     rewriter.create<memref::StoreOp>(loc, value, buffer, size);
 
     // Increment the size of the buffer by 1.
     Value c1 = constantIndex(rewriter, loc, 1);
     size = rewriter.create<arith::AddIOp>(loc, size, c1);
     rewriter.create<memref::StoreOp>(loc, size, bufferSizes, idx);
 
     rewriter.replaceOp(op, buffer);

mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir

 //       CHECK: %[[P2:.*]] = arith.addi %[[P]], %[[C1]]
 //       CHECK: memref.store %[[P2]], %[[A]]{{\[}}%[[C2]]]
 //       CHECK: return %[[M]] : memref<?xf64>
 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>
 }
 
+// CHECK-LABEL: func @sparse_push_back_inbound(
+//  CHECK-SAME: %[[A:.*]]: memref<?xindex>,
+//  CHECK-SAME: %[[B:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[C:.*]]: f64) -> memref<?xf64> {
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//       CHECK: %[[P:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
+//       CHECK: memref.store %[[C]], %[[B]]{{\[}}%[[P]]]
+//       CHECK: %[[P2:.*]] = arith.addi %[[P]], %[[C1]]
+//       CHECK: memref.store %[[P2]], %[[A]]{{\[}}%[[C2]]]
+//       CHECK: return %[[B]] : memref<?xf64>
+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 to memref<?xf64>
+  return %0 : memref<?xf64>
+}
+
 // CHECK-LABEL:   func.func private @_sparse_less_than_1_i8(
 // CHECK-SAME:                                              %[[I:arg0]]: index,
 // CHECK-SAME:                                              %[[J:.*]]: index,
 // CHECK-SAME:                                              %[[X0:.*]]: memref<?xi8>) -> i1 {
 // CHECK:           %[[VI:.*]] = memref.load %[[X0]]{{\[}}%[[I]]]
 // CHECK:           %[[VJ:.*]] = memref.load %[[X0]]{{\[}}%[[J]]]
 // CHECK:           %[[C:.*]] = arith.cmpi ult, %[[VI]], %[[VJ]]
 // CHECK:           return %[[C]]

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       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>
 }
 
 // -----
 
+// CHECK-LABEL: func @sparse_push_back_inbound(
+//  CHECK-SAME: %[[A:.*]]: memref<?xindex>,
+//  CHECK-SAME: %[[B:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[C:.*]]: f64) -> memref<?xf64> {
+//       CHECK: %[[D:.*]] = sparse_tensor.push_back inbounds %[[A]], %[[B]], %[[C]] {idx = 2 : index} : memref<?xindex>, memref<?xf64>, f64 to memref<?xf64>
+//       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 to memref<?xf64>
+  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