[mlir][sparse] Introducing options for the SparseTensorConversion pass

[wrengr]

Summary

This is work towards: https://github.com/llvm/llvm-project/issues/51652

This differential sets up the options and threads them through everywhere, but doesn't actually use them yet. The differential that finally makes use of them is D122061, which is the final differential in the chain that fixes bug 51652.

[aartbik]

I have some comments, but okay to go after you address these.

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

93

Note that in general, I have followed the approach in which we completely separate the "mechanism" (the how to apply a transformation) from "policy" (the when and where to apply, i.e. the heuristics). Here you sort of step into both worlds where kAuto is entering the policy world, while the kVia/kDirect are still just two mechanisms. So I would be okay with just the two mechamisms for now (although you explained offline why kAuto is there). If we keep all three, perhaps mention a bit on policy vs mechanism (in general in this file perhaps at top, and concretely here).

And yes, I know, I am deviating from my regular "terse" commenting path ;-) ;-)

96

Good to add this reference, but perhaps push it to the actual implementation part, since there is not much background needed for the Via vs Direct enum here

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

467

is the underscore name really neeed? can't we just call it tpConv or so if you avoid a name clash?

[wrengr]

Adding commentary about separating mechanism vs policy.

[wrengr]

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

96

I added the reference just to give a bit more explanation about what "direct" means, since there are many possible options for direct conversion in the literal/colloquial sense (e.g., the AMT approach I mentioned offline). Honestly, I don't really like the name "direct" because of this ambiguity; but then I couldn't think of a more explicit name (other than something like alaTACO, which also isn't great).

...Maybe a name like kViaNNZ would be better, since it uses the SparseTensorNNZ class as the intermediate? (Not that I'm especially fond of the name of that class either.) That would make room for adding kViaAMT (etc) down the road

[wrengr]

renaming variable to avoid underscore

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

struct SparseCompilerOptions
 
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
     return SparsificationOptions(sparseParallelizationStrategy(parallelization),
                                  sparseVectorizationStrategy(vectorization),
                                  vectorLength, enableSIMDIndex32);
   }
 
+  // These options must be kept in sync with `SparseTensorConversionBase`.
+  PassOptions::Option<int32_t> sparseToSparse{
+      *this, "s2s-strategy",
+      desc("Set the strategy for sparse-to-sparse conversion"), init(0)};
+
+  /// Projects out the options for `createSparsificationPass`.
+  SparseTensorConversionOptions sparseTensorConversionOptions() const {
+    return SparseTensorConversionOptions(
+        sparseToSparseConversionStrategy(sparseToSparse));
+  }
+
   // These options must be kept in sync with `ConvertVectorToLLVMBase`.
   // TODO(wrengr): does `indexOptimizations` differ from `enableSIMDIndex32`?
   PassOptions::Option<bool> reassociateFPReductions{
       *this, "reassociate-fp-reductions",
       desc("Allows llvm to reassociate floating-point reductions for speed"),
       init(false)};
   PassOptions::Option<bool> indexOptimizations{
       *this, "enable-index-optimizations",

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

 //===- Passes.h - Sparse tensor pass entry points ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines prototypes of all sparse tensor passes.
 //
+// In general, this file takes the approach of keeping "mechanism" (the
+// actual steps of applying a transformation) completely separate from
+// "policy" (heuristics for when and where to apply transformations).
+// The only exception is in `SparseToSparseConversionStrategy`; for which,
+// see further discussion there.
+//
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
 
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 // Forward.
 class TypeConverter;
 
+//===----------------------------------------------------------------------===//
+// The Sparsification pass.
+//===----------------------------------------------------------------------===//
+
 /// Defines a parallelization strategy. Any independent loop is a candidate
 /// for parallelization. The loop is made parallel if (1) allowed by the
 /// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse
 /// outermost loop only), and (2) the generated code is an actual for-loop
 /// (and not a co-iterating while-loop).
 enum class SparseParallelizationStrategy {
   kNone,
   kDenseOuterLoop,
enum class SparseVectorizationStrategy {
   kNone,
   kDenseInnerLoop,
   kAnyStorageInnerLoop
 };
 
 /// Converts command-line vectorization flag to the strategy enum.
 SparseVectorizationStrategy sparseVectorizationStrategy(int32_t flag);
 
-/// Sparsification options.
+/// Options for the Sparsification pass.
 struct SparsificationOptions {
   SparsificationOptions(SparseParallelizationStrategy p,
                         SparseVectorizationStrategy v, unsigned vl, bool e)
       : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
         enableSIMDIndex32(e) {}
   SparsificationOptions()
       : SparsificationOptions(SparseParallelizationStrategy::kNone,
                               SparseVectorizationStrategy::kNone, 1u, false) {}
   SparseParallelizationStrategy parallelizationStrategy;
   SparseVectorizationStrategy vectorizationStrategy;
   unsigned vectorLength;
   bool enableSIMDIndex32;
 };
 
 /// Sets up sparsification rewriting rules with the given options.
 void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
-/// Sets up sparse tensor conversion rules.
-void populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
-                                            RewritePatternSet &patterns);
-
 std::unique_ptr<Pass> createSparsificationPass();
 std::unique_ptr<Pass>
 createSparsificationPass(const SparsificationOptions &options);
+
+//===----------------------------------------------------------------------===//
+// The SparseTensorConversion pass.
+//===----------------------------------------------------------------------===//
+
+/// Defines a strategy for implementing sparse-to-sparse conversion.
+/// `kAuto` leaves it up to the compiler to automatically determine
+/// the method used.  `kViaCOO` converts the source tensor to COO and
+/// then converts the COO to the target format.  `kDirect` converts
+/// directly via the algorithm in <https://arxiv.org/abs/2001.02609>;
+/// however, beware that there are many formats not supported by this
+/// conversion method.
+///
+/// The presence of the `kAuto` option violates our usual goal of keeping
+/// policy completely separated from mechanism.  The reason it exists is
+/// because (at present) this strategy can only be specified on a per-file
+/// basis.  To see why this is a problem, note that `kDirect` cannot
+/// support certain conversions; so if there is no `kAuto` setting,
+/// then whenever a file contains a single non-`kDirect`-able conversion
+/// the user would be forced to use `kViaCOO` for all conversions in
+/// that file!  In the future, instead of using this enum as a `Pass`
+/// option, we could instead move it to being an attribute on the
+/// conversion op; at which point `kAuto` would no longer be necessary.
+enum class SparseToSparseConversionStrategy { kAuto, kViaCOO, kDirect };
+
+/// Converts command-line sparse2sparse flag to the strategy enum.
+SparseToSparseConversionStrategy sparseToSparseConversionStrategy(int32_t flag);
+
+/// SparseTensorConversion options.
+struct SparseTensorConversionOptions {
+  SparseTensorConversionOptions(SparseToSparseConversionStrategy s2s)
+      : sparseToSparseStrategy(s2s) {}
+  SparseTensorConversionOptions()
+      : SparseTensorConversionOptions(SparseToSparseConversionStrategy::kAuto) {
+  }
+  SparseToSparseConversionStrategy sparseToSparseStrategy;
+};
+
+/// Sets up sparse tensor conversion rules.
+void populateSparseTensorConversionPatterns(
+    TypeConverter &typeConverter, RewritePatternSet &patterns,
+    const SparseTensorConversionOptions &options =
+        SparseTensorConversionOptions());
+
 std::unique_ptr<Pass> createSparseTensorConversionPass();
+std::unique_ptr<Pass>
+createSparseTensorConversionPass(const SparseTensorConversionOptions &options);
 
 //===----------------------------------------------------------------------===//
 // 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

let dependentDialects = [
     "bufferization::BufferizationDialect",
     "LLVM::LLVMDialect",
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
     "vector::VectorDialect",
   ];
+  let options = [
+    Option<"sparseToSparse", "s2s-strategy", "int32_t", "0",
+           "Set the strategy for sparse-to-sparse conversion">,
+  ];
 }
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

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

 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   // TODO(wrengr): ensure the original `pm` is for ModuleOp
   pm.addNestedPass<FuncOp>(createLinalgGeneralizationPass());
   pm.addPass(createLinalgElementwiseOpFusionPass());
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
-  pm.addPass(createSparseTensorConversionPass());
+  pm.addPass(createSparseTensorConversionPass(
+      options.sparseTensorConversionOptions()));
   pm.addNestedPass<FuncOp>(createLinalgBufferizePass());
   pm.addNestedPass<FuncOp>(vector::createVectorBufferizePass());
   pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<FuncOp>(createConvertVectorToSCFPass());
   pm.addNestedPass<FuncOp>(createConvertSCFToCFPass());
   pm.addPass(func::createFuncBufferizePass());
   pm.addPass(arith::createConstantBufferizePass());
   pm.addNestedPass<FuncOp>(createTensorBufferizePass());

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

newParams(rewriter, params, op, stp, enc, Action::kEmpty,
               adaptor.getOperands());
     rewriter.replaceOp(op, genNewCall(rewriter, op, params));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
+  /// Options to control sparse code generation.
+  SparseTensorConversionOptions options;
+
+public:
   using OpConversionPattern::OpConversionPattern;
+  SparseTensorConvertConverter(MLIRContext *context,
+                               SparseTensorConversionOptions o)
+      : OpConversionPattern<ConvertOp>(context), options(o) {}
+  SparseTensorConvertConverter(TypeConverter &typeConv, MLIRContext *context,
+                               SparseTensorConversionOptions o)
+      : OpConversionPattern<ConvertOp>(typeConv, context), options(o) {}
+
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     Type resType = op.getType();
     Type srcType = op.source().getType();
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(srcType);

 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
-void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
-                                                  RewritePatternSet &patterns) {
+void mlir::populateSparseTensorConversionPatterns(
+    TypeConverter &typeConverter, RewritePatternSet &patterns,
+    const SparseTensorConversionOptions &options) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
                SparseCastConverter, SparseTensorNewConverter,
-               SparseTensorInitConverter, SparseTensorConvertConverter,
-               SparseTensorReleaseConverter, SparseTensorToPointersConverter,
-               SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
-               SparseTensorLoadConverter, SparseTensorLexInsertConverter,
-               SparseTensorExpandConverter, SparseTensorCompressConverter,
-               SparseTensorOutConverter>(typeConverter, patterns.getContext());
+               SparseTensorInitConverter, SparseTensorReleaseConverter,
+               SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
+               SparseTensorToValuesConverter, SparseTensorLoadConverter,
+               SparseTensorLexInsertConverter, SparseTensorExpandConverter,
+               SparseTensorCompressConverter, SparseTensorOutConverter>(
+      typeConverter, patterns.getContext());
+  patterns.add<SparseTensorConvertConverter>(typeConverter,
+                                             patterns.getContext(), options);
 }

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

static Optional<Type> convertSparseTensorTypes(Type type) {
     if (getSparseTensorEncoding(type) != nullptr)
       return LLVM::LLVMPointerType::get(IntegerType::get(type.getContext(), 8));
     return llvm::None;
   }
 };
 
 struct SparseTensorConversionPass
     : public SparseTensorConversionBase<SparseTensorConversionPass> {
+
+  SparseTensorConversionPass() = default;
+  SparseTensorConversionPass(const SparseTensorConversionPass &pass) = default;
+  SparseTensorConversionPass(const SparseTensorConversionOptions &options) {
+    sparseToSparse = static_cast<int32_t>(options.sparseToSparseStrategy);
+  }
+
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeConverter converter;
     ConversionTarget target(*ctx);
     // Everything in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
     // All dynamic rules below accept new function, call, return, and tensor
void runOnOperation() override {
     // The following operations and dialects may be introduced by the
     // rewriting rules, and are therefore marked as legal.
     target.addLegalOp<arith::CmpFOp, arith::CmpIOp, arith::ConstantOp,
                       arith::IndexCastOp, linalg::FillOp, linalg::YieldOp,
                       tensor::ExtractOp>();
     target
         .addLegalDialect<bufferization::BufferizationDialect, LLVM::LLVMDialect,
                          memref::MemRefDialect, scf::SCFDialect>();
+    // Translate strategy flags to strategy options.
+    SparseTensorConversionOptions options(
+        sparseToSparseConversionStrategy(sparseToSparse));
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<FuncOp>(patterns,
                                                              converter);
     populateCallOpTypeConversionPattern(patterns, converter);
-    populateSparseTensorConversionPatterns(converter, patterns);
+    populateSparseTensorConversionPatterns(converter, patterns, options);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
default:
     return SparseVectorizationStrategy::kNone;
   case 1:
     return SparseVectorizationStrategy::kDenseInnerLoop;
   case 2:
     return SparseVectorizationStrategy::kAnyStorageInnerLoop;
   }
 }
 
+SparseToSparseConversionStrategy
+mlir::sparseToSparseConversionStrategy(int32_t flag) {
+  switch (flag) {
+  default:
+    return SparseToSparseConversionStrategy::kAuto;
+  case 1:
+    return SparseToSparseConversionStrategy::kViaCOO;
+  case 2:
+    return SparseToSparseConversionStrategy::kDirect;
+  }
+}
+
 std::unique_ptr<Pass> mlir::createSparsificationPass() {
   return std::make_unique<SparsificationPass>();
 }
 
 std::unique_ptr<Pass>
 mlir::createSparsificationPass(const SparsificationOptions &options) {
   return std::make_unique<SparsificationPass>(options);
 }
 
 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);
+}