[mlir][sparse] Add option enable-buffer-initialization to initialize the memory buffers for sparse tensors to support debugging.

[aartbik]

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

711

as discussed offline, no env variable but a compiler option

[bixia]

Use compiler option instead of env variable.

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

PassOptions::Option<bool> enableRuntimeLibrary{
       desc("Enable runtime library for manipulating sparse tensors"),
       // TODO: Disable runtime library by default after feature complete.
       init(true)};
 
   PassOptions::Option<bool> testBufferizationAnalysisOnly{
       *this, "test-bufferization-analysis-only",
       desc("Run only the inplacability analysis"), init(false)};
 
+  PassOptions::Option<bool> enableBufferInitialization{
+      *this, "enable-buffer-initialization",
+      desc("Enable zero-initialization of memory buffers"), init(false)};
+
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
     return SparsificationOptions(parallelization);
   }
 
   // These options must be kept in sync with `SparseTensorConversionBase`.
   PassOptions::Option<int32_t> sparseToSparse{
       *this, "s2s-strategy",

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

 
 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> createDenseBufferizationPass(
     const bufferization::OneShotBufferizationOptions &options);
 
-void populateSparseBufferRewriting(RewritePatternSet &patterns);
-std::unique_ptr<Pass> createSparseBufferRewritePass();
+void populateSparseBufferRewriting(RewritePatternSet &patterns,
+                                   bool enableBufferInitialization);
+std::unique_ptr<Pass>
+createSparseBufferRewritePass(bool enableBufferInitialization = false);
 
 //===----------------------------------------------------------------------===//
 // Registration.
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_

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

def SparseBufferRewrite : Pass<"sparse-buffer-rewrite", "ModuleOp"> {
   let constructor = "mlir::createSparseBufferRewritePass()";
   let dependentDialects = [
     "arith::ArithDialect",
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
+  let options = [
+    Option<"enableBufferInitialization", "enable-buffer-initialization", "bool",
+           "false", "Enable zero-initialization of the memory buffers">,
+  ];
 }
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

if (options.testBufferizationAnalysisOnly)
     return;
   pm.addPass(createSparseTensorRewritePass(options.enableRuntimeLibrary));
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
   if (options.enableRuntimeLibrary)
     pm.addPass(createSparseTensorConversionPass(
         options.sparseTensorConversionOptions()));
   else
     pm.addPass(createSparseTensorCodegenPass());
-  pm.addPass(createSparseBufferRewritePass());
+  pm.addPass(createSparseBufferRewritePass(options.enableBufferInitialization));
   pm.addPass(createDenseBufferizationPass(
       getBufferizationOptions(/*analysisOnly=*/false)));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());

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

 //===---------------------------------------------------------------------===//
 
 namespace {
 
 /// Sparse rewriting rule for the push_back operator.
 struct PushBackRewriter : OpRewritePattern<PushBackOp> {
 public:
   using OpRewritePattern<PushBackOp>::OpRewritePattern;
+  PushBackRewriter(MLIRContext *context, bool enableInit)
+      : OpRewritePattern(context), enableBufferInitialization(enableInit) {}
   LogicalResult matchAndRewrite(PushBackOp op,
                                 PatternRewriter &rewriter) const override {
     // 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)
if (!op.getInbounds()) {
         rewriter.create<scf::YieldOp>(loc, after->getArguments());
 
         rewriter.setInsertionPointAfter(whileOp);
         capacity = whileOp.getResult(0);
       }
 
       Value newBuffer =
           rewriter.create<memref::ReallocOp>(loc, bufferType, buffer, capacity);
+      if (enableBufferInitialization) {
+        Value fillSize = rewriter.create<arith::SubIOp>(loc, capacity, newSize);
+        Value fillValue = rewriter.create<arith::ConstantOp>(
+            loc, value.getType(), rewriter.getZeroAttr(value.getType()));
+        Value subBuffer = rewriter.create<memref::SubViewOp>(
+            loc, newBuffer, /*offset=*/ValueRange{newSize},
+            /*size=*/ValueRange{fillSize},
+            /*step=*/ValueRange{constantIndex(rewriter, loc, 1)});
+        rewriter.create<linalg::FillOp>(loc, fillValue, subBuffer);
+      }
       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);
if (nIsOne) {
       rewriter.create<linalg::FillOp>(loc, value, subBuffer);
     }
 
     // Update the buffer size.
     rewriter.create<memref::StoreOp>(loc, newSize, bufferSizes, idx);
     rewriter.replaceOp(op, buffer);
     return success();
   }
+
+private:
+  bool enableBufferInitialization;
 };
 
 /// Sparse rewriting rule for the sort operator.
 struct SortRewriter : public OpRewritePattern<SortOp> {
 public:
   using OpRewritePattern<SortOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SortOp op,

 };
 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
-void mlir::populateSparseBufferRewriting(RewritePatternSet &patterns) {
-  patterns.add<PushBackRewriter, SortRewriter>(patterns.getContext());
+void mlir::populateSparseBufferRewriting(RewritePatternSet &patterns,
+                                         bool enableBufferInitialization) {
+  patterns.add<PushBackRewriter>(patterns.getContext(),
+                                 enableBufferInitialization);
+  patterns.add<SortRewriter>(patterns.getContext());
 }

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

struct SparseTensorCodegenPass
   }
 };
 
 struct SparseBufferRewritePass
     : public impl::SparseBufferRewriteBase<SparseBufferRewritePass> {
 
   SparseBufferRewritePass() = default;
   SparseBufferRewritePass(const SparseBufferRewritePass &pass) = default;
+  SparseBufferRewritePass(bool enableInit) {
+    enableBufferInitialization = enableInit;
+  }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    populateSparseBufferRewriting(patterns);
+    populateSparseBufferRewriting(patterns, enableBufferInitialization);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Strategy flag methods.
std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(
     const SparseTensorConversionOptions &options) {
   return std::make_unique<SparseTensorConversionPass>(options);
 }
 
 std::unique_ptr<Pass> mlir::createSparseTensorCodegenPass() {
   return std::make_unique<SparseTensorCodegenPass>();
 }
 
-std::unique_ptr<Pass> mlir::createSparseBufferRewritePass() {
-  return std::make_unique<SparseBufferRewritePass>();
+std::unique_ptr<Pass>
+mlir::createSparseBufferRewritePass(bool enableBufferInitialization) {
+  return std::make_unique<SparseBufferRewritePass>(enableBufferInitialization);
 }