[mlir][sparse] Add option enable-buffer-initialization to the sparse-tensor-codegen pass.

[aartbik]

Can we add a check test for this. Preferably just a separate one (I mean not intermixed with codegen.mlir)?

[bixia]

Add a FileCheck test.

[aartbik]

mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir

16

indentation of 3x CHECK on linalg.fill?

32

newline?

[bixia]

Fix test file format.

[bixia]

Fix test.

[bixia]

Fix the constructor.

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

 /// Sparse tensor type converter into an actual buffer.
 class SparseTensorTypeToBufferConverter : public TypeConverter {
 public:
   SparseTensorTypeToBufferConverter();
 };
 
 /// Sets up sparse tensor conversion rules.
 void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
-                                         RewritePatternSet &patterns);
+                                         RewritePatternSet &patterns,
+                                         bool enableBufferInitialization);
 
-std::unique_ptr<Pass> createSparseTensorCodegenPass();
+std::unique_ptr<Pass>
+createSparseTensorCodegenPass(bool enableBufferInitialization = false);
 
 //===----------------------------------------------------------------------===//
 // The SparseTensorRewriting pass.
 //===----------------------------------------------------------------------===//
 
 void populateSparseTensorRewriting(RewritePatternSet &patterns, bool enableRT,
                                    bool enableForeach, bool enableConvert);
 

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

def SparseTensorCodegen : Pass<"sparse-tensor-codegen", "ModuleOp"> {
   let dependentDialects = [
     "arith::ArithDialect",
     "bufferization::BufferizationDialect",
     "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">,
+  ];
 }
 
 def SparseBufferRewrite : Pass<"sparse-buffer-rewrite", "ModuleOp"> {
   let summary = "Rewrite sparse primitives on buffers to actual code";
   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.

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

void mlir::sparse_tensor::buildSparseCompiler(
   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(
+        createSparseTensorCodegenPass(options.enableBufferInitialization));
   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.

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

static void allocSchemeForRank(OpBuilder &builder, Location loc,
   // Reached values array so prepare for an insertion.
   Value valZero = constantZero(builder, loc, rtp.getElementType());
   createPushback(builder, loc, fields, field, valZero, linear);
   assert(fields.size() == ++field);
 }
 
 /// Creates allocation operation.
 static Value createAllocation(OpBuilder &builder, Location loc, Type type,
-                              Value sz) {
+                              Value sz, bool enableInit) {
   auto memType = MemRefType::get({ShapedType::kDynamicSize}, type);
-  return builder.create<memref::AllocOp>(loc, memType, sz);
+  Value buffer = builder.create<memref::AllocOp>(loc, memType, sz);
+  if (enableInit) {
+    Value fillValue =
+        builder.create<arith::ConstantOp>(loc, type, builder.getZeroAttr(type));
+    builder.create<linalg::FillOp>(loc, fillValue, buffer);
+  }
+  return buffer;
 }
 
 /// Creates allocation for each field in sparse tensor type. Note that
 /// for all dynamic memrefs, the memory size is really the capacity of
 /// the "vector", while the actual size resides in the sizes array.
 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
 ///       on the required capacities (see heuristic variable).
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
-                              ValueRange dynSizes,
+                              ValueRange dynSizes, bool enableInit,
                               SmallVectorImpl<Value> &fields) {
   auto enc = getSparseTensorEncoding(type);
   assert(enc);
   // Construct the basic types.
   unsigned idxWidth = enc.getIndexBitWidth();
   unsigned ptrWidth = enc.getPointerBitWidth();
   RankedTensorType rtp = type.cast<RankedTensorType>();
   Type indexType = builder.getIndexType();
Value dimSizes =
       builder.create<memref::AllocOp>(loc, MemRefType::get({rank}, indexType));
   Value memSizes = builder.create<memref::AllocOp>(
       loc, MemRefType::get({getMemSizesIndex(lastField)}, indexType));
   fields.push_back(dimSizes);
   fields.push_back(memSizes);
   // Per-dimension storage.
   for (unsigned r = 0; r < rank; r++) {
     if (isCompressedDim(rtp, r)) {
-      fields.push_back(createAllocation(builder, loc, ptrType, heuristic));
-      fields.push_back(createAllocation(builder, loc, idxType, heuristic));
+      fields.push_back(
+          createAllocation(builder, loc, ptrType, heuristic, enableInit));
+      fields.push_back(
+          createAllocation(builder, loc, idxType, heuristic, enableInit));
     } else if (isSingletonDim(rtp, r)) {
-      fields.push_back(createAllocation(builder, loc, idxType, heuristic));
+      fields.push_back(
+          createAllocation(builder, loc, idxType, heuristic, enableInit));
     } else {
       assert(isDenseDim(rtp, r)); // no fields
     }
   }
   // The values array.
-  fields.push_back(createAllocation(builder, loc, eltType, heuristic));
+  fields.push_back(
+      createAllocation(builder, loc, eltType, heuristic, enableInit));
   assert(fields.size() == lastField);
   // Initialize the storage scheme to an empty tensor. Initialized memSizes
   // to all zeros, sets the dimSizes to known values and gives all pointer
   // fields an initial zero entry, so that it is easier to maintain the
   // "linear + 1" length property.
   builder.create<linalg::FillOp>(
       loc, ValueRange{constantZero(builder, loc, indexType)},
       ValueRange{memSizes}); // zero memSizes
public:
   }
 };
 
 /// Sparse codgen rule for the alloc operator.
 class SparseTensorAllocConverter
     : public OpConversionPattern<bufferization::AllocTensorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
+  SparseTensorAllocConverter(MLIRContext *context, bool enableInit)
+      : OpConversionPattern(context), enableBufferInitialization(enableInit) {}
   LogicalResult
   matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     RankedTensorType resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     if (op.getCopy())
       return rewriter.notifyMatchFailure(op, "tensor copy not implemented");
 
     // Construct allocation for each field.
     Location loc = op.getLoc();
     SmallVector<Value, 8> fields;
-    createAllocFields(rewriter, loc, resType, adaptor.getOperands(), fields);
+    createAllocFields(rewriter, loc, resType, adaptor.getOperands(),
+                      enableBufferInitialization, fields);
     // Replace operation with resulting memrefs.
     rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields));
     return success();
   }
+
+private:
+  bool enableBufferInitialization;
 };
 
 /// Sparse codegen rule for the dealloc operator.
 class SparseTensorDeallocConverter
     : public OpConversionPattern<bufferization::DeallocTensorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult

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

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

struct SparseTensorConversionPass
   }
 };
 
 struct SparseTensorCodegenPass
     : public impl::SparseTensorCodegenBase<SparseTensorCodegenPass> {
 
   SparseTensorCodegenPass() = default;
   SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;
+  SparseTensorCodegenPass(bool enableInit) {
+    enableBufferInitialization = enableInit;
+  }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeToBufferConverter converter;
     ConversionTarget target(*ctx);
     // Most ops in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
target.addLegalDialect<arith::ArithDialect,
                            bufferization::BufferizationDialect,
                            memref::MemRefDialect, scf::SCFDialect>();
     target.addLegalOp<UnrealizedConversionCastOp>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
-    populateSparseTensorCodegenPatterns(converter, patterns);
+    populateSparseTensorCodegenPatterns(converter, patterns,
+                                        enableBufferInitialization);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 struct SparseBufferRewritePass
     : public impl::SparseBufferRewriteBase<SparseBufferRewritePass> {
std::unique_ptr<Pass> mlir::createSparseTensorConversionPass() {
   return std::make_unique<SparseTensorConversionPass>();
 }
 
 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::createSparseTensorCodegenPass(bool enableBufferInitialization) {
+  return std::make_unique<SparseTensorCodegenPass>(enableBufferInitialization);
 }
 
 std::unique_ptr<Pass>
 mlir::createSparseBufferRewritePass(bool enableBufferInitialization) {
   return std::make_unique<SparseBufferRewritePass>(enableBufferInitialization);
 }

mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir

+// RUN: mlir-opt %s --sparse-tensor-codegen=enable-buffer-initialization=true  --canonicalize --cse | FileCheck %s
+
+#SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
+
+// CHECK-LABEL: func @sparse_alloc_sparse_vector(
+//  CHECK-SAME: %[[A:.*]]: index) ->
+//  CHECK-SAME: memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.{{0*}}e+00 : f64
+//       CHECK: %[[T0:.*]] = memref.alloc() : memref<1xindex>
+//       CHECK: %[[T1:.*]] = memref.alloc() : memref<3xindex>
+//       CHECK: %[[T2:.*]] = memref.alloc() : memref<16xindex>
+//       CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<16xindex> to memref<?xindex>
+//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T2]] : memref<16xindex>)
+//       CHECK: %[[T4:.*]] = memref.alloc() : memref<16xindex>
+//       CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<16xindex> to memref<?xindex>
+//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T4]] : memref<16xindex>)
+//       CHECK: %[[T6:.*]] = memref.alloc() : memref<16xf64>
+//       CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<16xf64> to memref<?xf64>
+//       CHECK: linalg.fill ins(%[[F0]] : f64) outs(%[[T6]] : memref<16xf64>)
+//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T1]] : memref<3xindex>)
+//       CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<1xindex>
+//       CHECK: %[[P0:.*]] = sparse_tensor.push_back %[[T1]], %[[T3]]
+//       CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[T1]], %[[P0]]
+//       CHECK: return %[[T0]], %[[T1]], %[[P1]], %[[T5]], %[[T7]] :
+func.func @sparse_alloc_sparse_vector(%arg0: index) -> tensor<?xf64, #SV> {
+  %0 = bufferization.alloc_tensor(%arg0) : tensor<?xf64, #SV>
+  %1 = sparse_tensor.load %0 : tensor<?xf64, #SV>
+  return %1 : tensor<?xf64, #SV>
+}