[mlir][sparse] Add codegen rule for the push_back operator.

[Peiming]

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

636

I am sorry that I missed your previous patch (which introduces the operator).

But would it be better if we get rid of idx from the push_back operator? the memref<?xindex> size array can first be viewed into a single memref<index> to construct the push_back operator? I feel this way might be cleaner?

Other than that, the patch LGTM.

[bixia]

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

636

Per offline discussion, we will re-evaluate this approach later as it introduces memory aliasing.

[aartbik]

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

20

Thanks for adding a integration!

[rriddle]

sparse_codegen_push_back.mlir is seemingly failing in asan builds:

==62819==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x602000011c34 at pc 0x55fd675ed144 bp 0x7ffd7761ed30 sp 0x7ffd7761e4f8
READ of size 8 at 0x602000011c34 thread T0
    #0 0x55fd675ed143 in memcpy (/home/river/llvm-project/build/bin/mlir-cpu-runner+0x4423143) (BuildId: 3d88374a90e38509)
    #1 0x7f104d2840b3  (<unknown module>)
    #2 0x7f104d28418c  (<unknown module>)
    #3 0x55fd68c4f97b in compileAndExecute((anonymous namespace)::Options&, mlir::ModuleOp, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig, void**) /home/river/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:234:3
    #4 0x55fd68c4daca in compileAndExecuteVoidFunction((anonymous namespace)::Options&, mlir::ModuleOp, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig) /home/river/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:250:10
    #5 0x55fd68c47dfa in mlir::JitRunnerMain(int, char**, mlir::DialectRegistry const&, mlir::JitRunnerConfig) /home/river/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:383:23
    #6 0x55fd67692d31 in main /home/river/llvm-project/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp:33:10
    #7 0x7f104f319082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 1878e6b475720c7c51969e69ab2d276fae6d1dee)
    #8 0x55fd675d2efd in _start (/home/river/llvm-project/build/bin/mlir-cpu-runner+0x4408efd) (BuildId: 3d88374a90e38509)

0x602000011c34 is located 0 bytes to the right of 4-byte region [0x602000011c30,0x602000011c34)
allocated by thread T0 here:
    #0 0x55fd6765560e in malloc (/home/river/llvm-project/build/bin/mlir-cpu-runner+0x448b60e) (BuildId: 3d88374a90e38509)
    #1 0x7f104d28402a  (<unknown module>)
    #2 0x55fd68c4daca in compileAndExecuteVoidFunction((anonymous namespace)::Options&, mlir::ModuleOp, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig) /home/river/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:250:10
    #3 0x55fd68c47dfa in mlir::JitRunnerMain(int, char**, mlir::DialectRegistry const&, mlir::JitRunnerConfig) /home/river/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:383:23
    #4 0x55fd67692d31 in main /home/river/llvm-project/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp:33:10
    #5 0x7f104f319082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 1878e6b475720c7c51969e69ab2d276fae6d1dee)

SUMMARY: AddressSanitizer: heap-buffer-overflow (/home/river/llvm-project/build/bin/mlir-cpu-runner+0x4423143) (BuildId: 3d88374a90e38509) in memcpy

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

matchAndRewrite(CompressOp op, OpAdaptor adaptor,
     rewriter.create<memref::DeallocOp>(loc, values);
     rewriter.create<memref::DeallocOp>(loc, filled);
     rewriter.create<memref::DeallocOp>(loc, added);
     rewriter.eraseOp(op);
     return success();
   }
 };
 
+/// Sparse codegen rule for the push_back operator.
+class SparsePushBackConverter : public OpConversionPattern<PushBackOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(PushBackOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Lower 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)++
+    Location loc = op->getLoc();
+    Value c0 = constantIndex(rewriter, loc, 0);
+    Value buffer = adaptor.getInBuffer();
+    Value capacity = rewriter.create<memref::DimOp>(loc, buffer, c0);
+    Value idx = constantIndex(rewriter, loc, op.getIdx().getZExtValue());
+    Value bufferSizes = adaptor.getBufferSizes();
+    Value size = rewriter.create<memref::LoadOp>(loc, bufferSizes, idx);
+    Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge,
+                                                size, capacity);
+    Value value = adaptor.getValue();
+    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.
+    rewriter.setInsertionPointAfter(ifOp);
+    buffer = ifOp.getResult(0);
+    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);
+    return success();
+  }
+};
+
 /// Base class for getter-like operations, e.g., to_indices, to_pointers.
 template <typename SourceOp, typename Base>
 class SparseGetterOpConverter : public OpConversionPattern<SourceOp> {
 public:
   using OpAdaptor = typename SourceOp::Adaptor;
   using OpConversionPattern<SourceOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,

 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorLoadConverter,
                SparseExpandConverter, SparseCompressConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToValuesConverter>(typeConverter, patterns.getContext());
+               SparsePushBackConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToValuesConverter>(
+      typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

func.func @sparse_compression(%arg0: tensor<8x8xf64, #CSR>,
                               %arg2: memref<?xf64>,
                               %arg3: memref<?xi1>,
                               %arg4: memref<?xindex>,
                               %arg5: index) {
   sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
     : tensor<8x8xf64, #CSR>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
   return
 }
+
+// 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: %[[M:.*]] = scf.if %[[T]] -> (memref<?xf64>) {
+//       CHECK:  %[[P1:.*]] = arith.muli %[[S]], %[[C2]]
+//       CHECK:  %[[M2:.*]] = memref.realloc %[[B]](%[[P1]])
+//       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: 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>
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_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
+    %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 to memref<?xf32>
+    %buffer3 = sparse_tensor.push_back %bufferSizes, %buffer2, %d1 {idx=0 : index} : memref<?xindex>, memref<?xf32>, f32 to memref<?xf32>
+
+    // 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>
+
+    // Release the buffers.
+    memref.dealloc %bufferSizes : memref<?xindex>
+    memref.dealloc %buffer3 : memref<?xf32>
+    return
+  }
+}
+