[Mlir-commits] [mlir] [MLIR][SparseTensor] Added Sparse Outer Loop Ordering Strategy (PR #172198)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Govind Malasani (gmalasan)

<details>
<summary>Changes</summary>

This PR builds upon the infrastructure set up for Sparse Tensor Loop Ordering Heuristics (#<!-- -->154656) and the already existing Dense Outer loop ordering strategy (#<!-- -->160168).

---
Full diff: https://github.com/llvm/llvm-project/pull/172198.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h (+1) 
- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+3-1) 
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp (+43-11) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 419ecda80e9a5..40b37dc05e92e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -62,6 +62,7 @@ namespace sparse_tensor {
 enum class LoopOrderingStrategy : unsigned {
   kDefault,
   kDenseOuter,
+  kSparseOuter,
 };
 
 } // namespace sparse_tensor
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 0b8562e484f51..7ad54a0d2218d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -87,7 +87,9 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> {
          clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDefault, "default",
                     "Default strategy (eagerly selects last loop in topological sort)"),
          clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDenseOuter, "dense-outer",
-                    "Prefer dense, then compressed, then singleton dimensions outermost"))}]>,
+                    "Prefer dense, then compressed, then singleton dimensions outermost"),
+         clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kSparseOuter, "sparse-outer",
+                    "Prefer singleton, then compressed, then dense dimensions outermost"))}]>,
   ];
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 99048034b4f0c..a0180d228a36a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -82,11 +82,19 @@ inline static bool includesDenseOutput(SortMask mask) {
 
 /// Returns a sparsity rank for loop ordering: lower values indicate
 /// dimensions that should be placed in outer loops.
-/// 0 = Dense, 1 = Compressed, 2 = Singleton, 3 = Other/Unknown.
+/// When preferDenseOuter is true the ranking is
+///   0 = Dense, 1 = Compressed, 2 = Singleton, 3 = Other/Unknown.
+/// Otherwise
+///   0 = Singleton, 1 = Compressed, 2 = Dense, 3 = Other/Unknown.
 static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
-                                    ArrayRef<AffineMap> allMaps) {
-  // Start with highest rank.
-  unsigned minRank = 3;
+                                    ArrayRef<AffineMap> allMaps,
+                                    bool preferDenseOuter) {
+  const unsigned denseRank = preferDenseOuter ? 0 : 2;
+  const unsigned singletonRank = preferDenseOuter ? 2 : 0;
+  const unsigned compressedRank = 1;
+  const unsigned unknownRank = 3;
+
+  unsigned minRank = unknownRank;
 
   for (auto [tensor, map] : llvm::zip(allTensors, allMaps)) {
     // Check if this loop accesses this tensor.
@@ -105,19 +113,19 @@ static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
     if (loopAccessesTensor) {
       const auto enc = getSparseTensorEncoding(tensor.getType());
       if (!enc) {
-        // Dense tensor - lowest rank.
-        return 0;
+        // Dense tensor.
+        return denseRank;
       } else {
         // Sparse tensor - check the level type for this dimension.
         auto lvlTypes = enc.getLvlTypes();
         if (tensorDim < lvlTypes.size()) {
           auto lvlType = lvlTypes[tensorDim];
           if (isDenseLT(lvlType)) {
-            return 0; // Dense level.
+            return denseRank; // Dense level.
           } else if (isCompressedLT(lvlType)) {
-            minRank = std::min(minRank, 1u); // Compressed level.
+            minRank = std::min(minRank, compressedRank); // Compressed level.
           } else if (isSingletonLT(lvlType)) {
-            minRank = std::min(minRank, 2u); // Singleton level.
+            minRank = std::min(minRank, singletonRank); // Singleton level.
           }
         }
       }
@@ -164,10 +172,34 @@ AffineMap IterationGraphSorter::topoSort() {
 
       // Find loop with minimum (lowest) sparsity rank.
       unsigned minLoop = it[0];
-      unsigned minRank = getLoopSparsityRank(minLoop, allTensors, allMaps);
+      unsigned minRank =
+          getLoopSparsityRank(minLoop, allTensors, allMaps, true);
+
+      for (auto candidateLoop : it) {
+        unsigned rank =
+            getLoopSparsityRank(candidateLoop, allTensors, allMaps, true);
+        if (rank < minRank || (rank == minRank && candidateLoop < minLoop)) {
+          minLoop = candidateLoop;
+          minRank = rank;
+        }
+      }
+      src = minLoop;
+      break;
+    }
+    case sparse_tensor::LoopOrderingStrategy::kSparseOuter: {
+      // Prefer singleton, then compressed, then dense dimensions outermost.
+      SmallVector<Value> allTensors = ins;
+      allTensors.push_back(out);
+      SmallVector<AffineMap> allMaps = loop2InsLvl;
+      allMaps.push_back(loop2OutLvl);
+
+      unsigned minLoop = it[0];
+      unsigned minRank =
+          getLoopSparsityRank(minLoop, allTensors, allMaps, false);
 
       for (auto candidateLoop : it) {
-        unsigned rank = getLoopSparsityRank(candidateLoop, allTensors, allMaps);
+        unsigned rank =
+            getLoopSparsityRank(candidateLoop, allTensors, allMaps, false);
         if (rank < minRank || (rank == minRank && candidateLoop < minLoop)) {
           minLoop = candidateLoop;
           minRank = rank;

``````````

</details>


https://github.com/llvm/llvm-project/pull/172198