[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

545

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

548

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

548

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

530

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

530

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.

545

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...

544

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)
+    //    while new_size > capacity
+    //      new_capacity = capacity*2
     //    new_buffer = realloc(buffer, new_capacity)
     // buffer = new_buffer
-    // store(buffer, value)
-    // size(buffer)++
+    // subBuffer = subviewof(buffer)
+    // linalg.fill subBuffer value
+    //
+    // update the buffer size with new_size
     //
     // 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();
+    uint64_t n = 1;
+    if (auto nAttr = op.getNAttr()) {
+      n = nAttr.getInt();
+    }
+    Value cn = constantIndex(rewriter, loc, n);
+    Value newSize = rewriter.create<arith::AddIOp>(loc, size, cn);
 
     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 (n == 1) {
         capacity = rewriter.create<arith::MulIOp>(loc, capacity, c2);
+      } else {
+        //   do { new_capacity *= 2 } while (size > new_capacity)
+        SmallVector<Type, 1> types{capacity.getType()};
+        scf::WhileOp whileOp = rewriter.create<scf::WhileOp>(
+            loc, types, SmallVector<Value, 1>{capacity});
+        // The before-region of the WhileOp.
+        Block *before =
+            rewriter.createBlock(&whileOp.getBefore(), {}, types, {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(), {}, types, {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 (n == 1) {
       rewriter.create<memref::StoreOp>(loc, value, buffer, size);
+    } else {
+      Value subBuffer = rewriter.create<memref::SubViewOp>(
+          loc, buffer, /*offset=*/ValueRange{size}, /*size=*/ValueRange{cn},
+          /*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) -> memref<?xf64> {
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//       CHECK: %[[P1:.*]] = memref.dim %[[B]], %[[C0]]
+//       CHECK: %[[S1:.*]] = memref.load %[[A]]{{\[}}%[[C2]]]
+//       CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[C10]] : 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]]] [10] [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) -> memref<?xf64> {
+  %0 = sparse_tensor.push_back %arg0, %arg1, %arg2 {idx = 2 : index, n=10 : index} : memref<?xindex>, memref<?xf64>, f64
+  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

func.func @entry() {
     %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 {idx=0 : index, n=10 : index} : memref<?xindex>, memref<?xf32>, f32
 
-    // 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
   }
 }