[mlir][sparse] Implement the rewrite for sparse_tensor.push_back a value n times.

[bixia]

Fix comment.

[bixia]

Make n a runtime value.

[bixia]

Rebase.

[Peiming]

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

544

Isn't there any O(1) algorithm for finding the closest power of 2?

547

I think you can simply do ArrayRef<Type>(capacity.getType()) to avoid creating an extra SmallVector and copying value to it.

[Peiming]

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

547

Actually you can simply do rewriter.create<scf::WhileOp>(loc, capactiy.getType(), capacity); as ArrayRef can be implicitly converted from a single element.

[aartbik]

very neat, thanks for writing this so quickly!

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

504

I think you need to keep the "buffer" in the "capacity(buffer)", as before

512

I actually like your original, which can now be

size(buffer) += n

524

this is of course very minor, but can we avoid creating the constant altogether for the default case and just set nIsOne directly (but of course, keep the check for when it is a runtime value that happens to be a constant 1)?

[bixia]

Address review comments.

[bixia]

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

524

Need to create the constant n for computing the newSize and updating size(buffer).
The current implementation handles push_back without n and push_back with explicit n==1 the same way.

544

the O(1) alg is like this. Per offline discussion, keep the current approach for simplicity.

If (a<b) {
  a = (a << (countLeadingZero(a) - countLeadingZero(b)))
  If (a < b)
    a = a*2
}

[aartbik]

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

505

you are missing a

new_capacity = capacity(buffer)

before the while-loop (not that we actually run the comments ;-) but still...

543

Note that you already have the pseudo code above in the comments, so I would not repeat too much,
Just generate the while-loop to find new capacity would be enough, I think

[bixia]

Rebase.

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

def SparseBufferRewrite : Pass<"sparse-buffer-rewrite", "ModuleOp"> {
   let description = [{
     A pass that rewrites sparse primitives on buffers to the MLIR implementation
     of the primitives. For example, sparse_tensor.sort operator is implemented
     in this pass.
   }];
   let constructor = "mlir::createSparseBufferRewritePass()";
   let dependentDialects = [
     "arith::ArithDialect",
+    "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
 }
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

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

 // primitives with memref operands.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;

 namespace {
 
 /// Sparse rewriting rule for the push_back operator.
 struct PushBackRewriter : OpRewritePattern<PushBackOp> {
 public:
   using OpRewritePattern<PushBackOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(PushBackOp op,
                                 PatternRewriter &rewriter) const override {
-    // Rewrite push_back(buffer, value) to:
-    // if (size(buffer) >= capacity(buffer))
-    //    new_capacity = capacity(buffer)*2
+    // Rewrite push_back(buffer, value, n) to:
+    // new_size = size(buffer) + n
+    // if (new_size > capacity(buffer))
+    //    while new_size > new_capacity
+    //      new_capacity = new_capacity*2
     //    new_buffer = realloc(buffer, new_capacity)
     // buffer = new_buffer
-    // store(buffer, value)
-    // size(buffer)++
+    // subBuffer = subviewof(buffer)
+    // linalg.fill subBuffer value
+    //
+    // size(buffer) += n
     //
     // 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 value = op.getValue();
 
+    Value n = op.getN() ? op.getN() : constantIndex(rewriter, loc, 1);
+    Value newSize = rewriter.create<arith::AddIOp>(loc, size, n);
+    auto nValue = dyn_cast_or_null<arith::ConstantIndexOp>(n.getDefiningOp());
+    bool nIsOne = (nValue && nValue.value() == 1);
+
     if (!op.getInbounds()) {
       Value cond = rewriter.create<arith::CmpIOp>(
-          loc, arith::CmpIPredicate::uge, size, capacity);
+          loc, arith::CmpIPredicate::ugt, newSize, capacity);
 
+      Value c2 = constantIndex(rewriter, loc, 2);
       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);
+      if (nIsOne) {
         capacity = rewriter.create<arith::MulIOp>(loc, capacity, c2);
+      } else {
+        // Use a do-while loop to calculate the new capacity as follows:
+        //   do { new_capacity *= 2 } while (size > new_capacity)
+        scf::WhileOp whileOp =
+            rewriter.create<scf::WhileOp>(loc, capacity.getType(), capacity);
+
+        // The before-region of the WhileOp.
+        Block *before = rewriter.createBlock(&whileOp.getBefore(), {},
+                                             {capacity.getType()}, {loc});
+        rewriter.setInsertionPointToEnd(before);
+
+        capacity =
+            rewriter.create<arith::MulIOp>(loc, before->getArgument(0), c2);
+        cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ugt,
+                                              newSize, capacity);
+        rewriter.create<scf::ConditionOp>(loc, cond, ValueRange{capacity});
+        // The after-region of the WhileOp.
+        Block *after = rewriter.createBlock(&whileOp.getAfter(), {},
+                                            {capacity.getType()}, {loc});
+        rewriter.setInsertionPointToEnd(after);
+        rewriter.create<scf::YieldOp>(loc, after->getArguments());
+
+        rewriter.setInsertionPointAfter(whileOp);
+        capacity = whileOp.getResult(0);
+      }
+
       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);
 
       // 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.
+    if (nIsOne) {
       rewriter.create<memref::StoreOp>(loc, value, buffer, size);
+    } else {
+      Value subBuffer = rewriter.create<memref::SubViewOp>(
+          loc, buffer, /*offset=*/ValueRange{size}, /*size=*/ValueRange{n},
+          /*step=*/ValueRange{constantIndex(rewriter, loc, 1)});
+      rewriter.create<linalg::FillOp>(loc, value, subBuffer);
+    }
 
-    // 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);
-
+    // Update the buffer size.
+    rewriter.create<memref::StoreOp>(loc, newSize, bufferSizes, idx);
     rewriter.replaceOp(op, buffer);
     return success();
   }
 };
 
 /// Sparse rewriting rule for the sort operator.
 struct SortRewriter : public OpRewritePattern<SortOp> {
 public:

mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir

 // RUN: mlir-opt %s -split-input-file --sparse-buffer-rewrite  --canonicalize --cse | FileCheck %s
 
 // CHECK-LABEL: func @sparse_push_back(
 //  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-DAG: %[[C0:.*]] = arith.constant 0 : index
-//       CHECK: %[[S:.*]] = memref.dim %[[B]], %[[C0]]
-//       CHECK: %[[P:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
-//       CHECK: %[[T:.*]] = arith.cmpi uge, %[[P]], %[[S]]
+//       CHECK: %[[P1:.*]] = memref.dim %[[B]], %[[C0]]
+//       CHECK: %[[S1:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
+//       CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[C1]] : index
+//       CHECK: %[[T:.*]] = arith.cmpi ugt, %[[S2]], %[[P1]]
 //       CHECK: %[[M:.*]] = scf.if %[[T]] -> (memref<?xf64>) {
-//       CHECK:  %[[P1:.*]] = arith.muli %[[S]], %[[C2]]
-//       CHECK:  %[[M2:.*]] = memref.realloc %[[B]](%[[P1]])
+//       CHECK:  %[[P2:.*]] = arith.muli %[[P1]], %[[C2]]
+//       CHECK:  %[[M2:.*]] = memref.realloc %[[B]](%[[P2]])
 //       CHECK:  scf.yield %[[M2]] : memref<?xf64>
 //       CHECK: } else {
 //       CHECK:  scf.yield %[[B]] : memref<?xf64>
 //       CHECK: }
-//       CHECK: memref.store %[[C]], %[[M]]{{\[}}%[[P]]]
-//       CHECK: %[[P2:.*]] = arith.addi %[[P]], %[[C1]]
-//       CHECK: memref.store %[[P2]], %[[A]]{{\[}}%[[C2]]]
+//       CHECK: memref.store %[[C]], %[[M]]{{\[}}%[[S1]]]
+//       CHECK: memref.store %[[S2]], %[[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
   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-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//       CHECK: %[[P1:.*]] = memref.dim %[[B]], %[[C0]]
+//       CHECK: %[[S1:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
+//       CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[D]] : index
+//       CHECK: %[[T:.*]] = arith.cmpi ugt, %[[S2]], %[[P1]]
+//       CHECK: %[[M:.*]] = scf.if %[[T]] -> (memref<?xf64>) {
+//       CHECK:   %[[P2:.*]] = scf.while (%[[I:.*]] = %[[P1]]) : (index) -> index {
+//       CHECK:     %[[P3:.*]] = arith.muli %[[I]], %[[C2]] : index
+//       CHECK:     %[[T2:.*]] = arith.cmpi ugt, %[[S2]], %[[P3]] : index
+//       CHECK:     scf.condition(%[[T2]]) %[[P3]] : index
+//       CHECK:   } do {
+//       CHECK:     ^bb0(%[[I2:.*]]: index):
+//       CHECK:     scf.yield %[[I2]] : index
+//       CHECK:   }  
+//       CHECK:  %[[M2:.*]] = memref.realloc %[[B]](%[[P2]])
+//       CHECK:  scf.yield %[[M2]] : memref<?xf64>
+//       CHECK: } else {
+//       CHECK:  scf.yield %[[B]] : memref<?xf64>
+//       CHECK: }
+//       CHECK: %[[S:.*]] = memref.subview %[[M]]{{\[}}%[[S1]]] {{\[}}%[[D]]] [1]
+//       CHECK: linalg.fill ins(%[[C]] : f64) outs(%[[S]]
+//       CHECK: memref.store %[[S2]], %[[A]]{{\[}}%[[C2]]]
+//       CHECK: return %[[M]] : memref<?xf64>
+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>
+}
+
+// -----
+
 // 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: %[[S1:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
+//       CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[C1]]
+//       CHECK: memref.store %[[C]], %[[B]]{{\[}}%[[S1]]]
+//       CHECK: memref.store %[[S2]], %[[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
   return %0 : memref<?xf64>
 }
 
 // -----
 

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_push_back.mlir

 // RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 module {
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
+    %c10 = arith.constant 10 : index
     %d0 = arith.constant 0.0 : f32
     %d1 = arith.constant 1.0 : f32
     %d2 = arith.constant 2.0 : f32
 
     %bufferSizes = memref.alloc(%c1) : memref<?xindex>
     %buffer = memref.alloc(%c1) : memref<?xf32>
 
     memref.store %c0, %bufferSizes[%c0] : memref<?xindex>
     %buffer2 = sparse_tensor.push_back %bufferSizes, %buffer, %d2 {idx=0 : index} : memref<?xindex>, memref<?xf32>, f32
-    %buffer3 = sparse_tensor.push_back %bufferSizes, %buffer2, %d1 {idx=0 : index} : memref<?xindex>, memref<?xf32>, f32
+    %buffer3 = sparse_tensor.push_back %bufferSizes, %buffer2, %d1, %c10 {idx=0 : index} : memref<?xindex>, memref<?xf32>, f32, index
 
-    // CHECK: ( 2 )
-    %sizeValue = vector.transfer_read %bufferSizes[%c0], %c0: memref<?xindex>, vector<1xindex>
-    vector.print %sizeValue : vector<1xindex>
-
-    // CHECK ( 2, 1 )
-    %bufferValue = vector.transfer_read %buffer3[%c0], %d0: memref<?xf32>, vector<2xf32>
-    vector.print %bufferValue : vector<2xf32>
+    // CHECK: 16
+    %capacity = memref.dim %buffer3, %c0 : memref<?xf32>
+    vector.print %capacity : index
+
+    // CHECK: ( 11 )
+    %size = vector.transfer_read %bufferSizes[%c0], %c0: memref<?xindex>, vector<1xindex>
+    vector.print %size : vector<1xindex>
+
+    // CHECK (  2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )
+    %values = vector.transfer_read %buffer3[%c0], %d0: memref<?xf32>, vector<11xf32>
+    vector.print %values : vector<11xf32>
 
     // Release the buffers.
     memref.dealloc %bufferSizes : memref<?xindex>
     memref.dealloc %buffer3 : memref<?xf32>
     return
   }
 }