[mlir][sparse] add new sparisification option for dependent index reduction-based codegen

[aartbik]

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

58

I would move (experimental) to the end, so e.g.

Enable slice-based algorithm to generate affine indices on sparse inputs (experimental feature)

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

85

same

[Peiming]

address comment, rename terminologies

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

PassOptions::Option<mlir::SparseParallelizationStrategy> parallelization{
           clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
                      "dense-any-loop",
                      "Enable dense parallelization for any loop."),
           clEnumValN(
               mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
               "any-storage-any-loop",
               "Enable sparse parallelization for any storage and loop."))};
 
+  PassOptions::Option<bool> enableIndexReduction{
+      *this, "enable-index-reduction",
+      desc("Enable dependent index reduction based algorithm to handle "
+           "non-trivial index expressions on sparse inputs (experimental "
+           "features)"),
+      init(false)};
+
   PassOptions::Option<bool> enableRuntimeLibrary{
       *this, "enable-runtime-library",
       desc("Enable runtime library for manipulating sparse tensors"),
       init(true)};
 
   PassOptions::Option<bool> testBufferizationAnalysisOnly{
       *this, "test-bufferization-analysis-only",
       desc("Run only the inplacability analysis"), init(false)};
struct SparseCompilerOptions
   PassOptions::Option<bool> x86Vector{
       *this, "enable-x86vector",
       desc("Enables the use of X86Vector dialect while lowering the vector "
            "dialect"),
       init(false)};
 
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
-    return SparsificationOptions(parallelization);
+    return SparsificationOptions(parallelization, enableIndexReduction);
   }
 
   /// Projects out the options for `createSparseTensorConversionPass`.
   SparseTensorConversionOptions sparseTensorConversionOptions() const {
     return SparseTensorConversionOptions(
         sparseToSparseConversionStrategy(sparseToSparse));
   }
 

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

enum class SparseParallelizationStrategy {
   // TODO: support reduction parallelization too?
 };
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
-  SparsificationOptions(SparseParallelizationStrategy p)
-      : parallelizationStrategy(p) {}
+  SparsificationOptions(SparseParallelizationStrategy p, bool idxReduc)
+      : parallelizationStrategy(p), enableIndexReduction(idxReduc) {}
   SparsificationOptions()
-      : SparsificationOptions(SparseParallelizationStrategy::kNone) {}
+      : SparsificationOptions(SparseParallelizationStrategy::kNone, false) {}
   SparseParallelizationStrategy parallelizationStrategy;
+  bool enableIndexReduction;
 };
 
 /// Sets up sparsification rewriting rules with the given options.
 void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
 std::unique_ptr<Pass> createSparsificationPass();

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

let dependentDialects = [
     "LLVM::LLVMDialect",
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
   // TODO(57514): These enum options are duplicated in Passes.h.
   let options = [
+    Option<"enableIndexReduction", "enable-index-reduction", "bool",
+           "false",
+           "Enable dependent index reduction based algorithm to handle non-trivial index expressions on sparse inputs (experimental features)">,
     Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
            "mlir::SparseParallelizationStrategy::kNone",
            "Set the parallelization strategy", [{llvm::cl::values(
              clEnumValN(mlir::SparseParallelizationStrategy::kNone, "none",
                         "Turn off sparse parallelization."),
              clEnumValN(mlir::SparseParallelizationStrategy::kDenseOuterLoop,
                         "dense-outer-loop",
                         "Enable dense outer loop sparse parallelization."),

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

 
 struct SparsificationPass
     : public impl::SparsificationPassBase<SparsificationPass> {
 
   SparsificationPass() = default;
   SparsificationPass(const SparsificationPass &pass) = default;
   SparsificationPass(const SparsificationOptions &options) {
     parallelization = options.parallelizationStrategy;
+    enableIndexReduction = options.enableIndexReduction;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     // Translate strategy flags to strategy options.
-    SparsificationOptions options(parallelization);
+    SparsificationOptions options(parallelization, enableIndexReduction);
     // Apply sparsification and cleanup rewriting.
     RewritePatternSet patterns(ctx);
     populateSparsificationPatterns(patterns, options);
     scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 

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

static bool translateBitsToTidLvlPairs(
     SmallVectorImpl<Level> &lvls, SmallVectorImpl<TensorId> &affineTids,
     SmallVectorImpl<Level> &affineLvls, SmallVectorImpl<AffineExpr> &exps) {
   const BitVector &simple = env.lat(li).simple;
   const TensorId outTid = env.merger().getOutTensorID();
   const std::optional<Level> outLvl = env.merger().getLvl(outTid, ldx);
 
   unsigned numloopCond = 0;
   bool hasNonUnique = false;
-  env.merger().foreachTensorLoopId(
-      li, [&, ldx](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
+  env.merger().foreachTensorLoopId(li, [&, ldx](TensorLoopId b, TensorId tid,
+                                                std::optional<Level> lvl,
                                                 DimLevelType dlt) {
-        if (simple.test(b)) {
-          if (isUndefDLT(dlt)) {
-            // An undefined dlt in the lattices, we probably mean to
-            // iterate based on the level of output tensor.  E.g., this
-            // could be a synthetic tensor (for invariants and sparse
-            // output tensor).
-            // out[i][j] = invariant; or a broadcast
-            // out[i][j] = in[i] (j is undef for input)
-            tid = outTid;
-            lvl = outLvl;
-            // Skips invalid lvl (e.g., when this is a zero ranked tensor).
-            if (!lvl)
-              return;
-          }
-          hasNonUnique = !isUniqueDLT(dlt) || hasNonUnique;
-          tids.push_back(tid);
-          lvls.push_back(*lvl);
-          numloopCond++;
-        } else if (isDenseDLT(dlt)) {
-          tids.push_back(tid);
-          lvls.push_back(*lvl);
-        } else {
-          assert(isUndefDLT(dlt));
-          linalg::GenericOp op = env.op();
-          if (tid >= op.getNumDpsInputs())
-            // We only handle affine expression on input tensors (for now).
-            return;
-          OpOperand *operand = &op->getOpOperand(tid);
-          const auto stt = getSparseTensorType(operand->get());
-          // Non-annotated dense tensors requires no special handling.
-          if (!stt.hasEncoding())
-            return;
 
-          ArrayRef<AffineExpr> affines =
-              op.getMatchingIndexingMap(operand).getResults();
-          const Level lvlRank = stt.getLvlRank();
-          assert(affines.size() == static_cast<size_t>(lvlRank));
-          for (Level l = 0; l < lvlRank; l++) {
-            // FIXME: `toOrigDim` is deprecated.
-            AffineExpr exp = affines[toOrigDim(stt.getEncoding(), l)];
-            // Skip simple affine expression and non-dense levels (which
-            // have their own filter loop).
-            if (exp.isa<AffineDimExpr>() || !stt.isDenseLvl(l))
-              continue;
 
-            // Constant affine expression are handled in genLoop
-            if (!exp.isa<AffineConstantExpr>()) {
-              bool isAtLoop = false;
-              if (isInvariantAffine(env, exp, ldx, isAtLoop) && isAtLoop) {
-                // If the compound affine is invariant and we are right at the
-                // level. We need to generate the address according to the
-                // affine expression. This is also the best place we can do it
-                // to avoid putting it inside inner loops.
-                // NOTE: It assumes that the levels of the input tensor are
-                // initialized in order (and it is also currently guaranteed by
-                // computeIterationGraph), another more admissible approach
-                // might be accepting out-of-order access between consecutive
-                // dense levels.
-                affineTids.push_back(tid);
-                affineLvls.push_back(l);
-                exps.push_back(exp);
-              }
-            }
-          }
-        }
-      });
 
   if (isDenseDLT(env.dlt(outTid, ldx))) {
     // Note that we generate dense indices of the output tensor
     // unconditionally, since they may not appear in the lattice, but may be
     // needed for linearized env.
     tids.push_back(outTid);
     lvls.push_back(*outLvl);
   }
public:
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Only accept single output operations without affine index on sparse
     // output.
     if (op.getNumDpsInits() != 1 || hasCompoundAffineOnSparseOut(op))
       return failure();
 
+    if (options.enableIndexReduction)
+      llvm_unreachable("not yet implemented");
+
     // Sets up a code generation environment.
     const unsigned numTensors = op->getNumOperands();
     const unsigned numLoops = op.getNumLoops();
     const unsigned numFilterLoops = getNumCompoundAffineOnSparseLvls(op);
     CodegenEnv env(op, options, numTensors, numLoops, numFilterLoops);
 
     // Detects sparse annotations and translates the per-level sparsity
     // information for all tensors to loop indices in the kernel.