[MLIR][SparseTensor] Loop ordering strategy infrastructure (flag) #154656

gmalasan · 2025-08-21T01:17:37Z

As discussed before, this PR adds the basic infrastructure/boiler plate for loop ordering strategies to be implemented.

If this looks ok, I wanted to also mention some of the heuristics that I would implement next, if they sound reasonable to you guys:

Parallel first : prioritize parallel loops over reduction loops
Dense outer : prioritize the most dense loops first
Sparse outer : the opposite, potentially useful in some cases?

There is another that I am considering, stride/memory aware, which would prioritize loops with better stride patterns (like sequential or linear). Not sure how well this carries over to Sparse Tensor though. Are there any ideas/heuristics that I should definitely try to implement?

As we discussed, I will try to incrementally add heuristics. Sorry for the delay on my end, and thank you so much for the feedback!

github-actions · 2025-08-21T01:17:56Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2025-08-21T01:18:27Z

@llvm/pr-subscribers-mlir-sparse

Author: Govind Malasani (gmalasan)

Changes

As discussed before, this PR adds the basic infrastructure/boiler plate for loop ordering strategies to be implemented.

If this looks ok, I wanted to also mention some of the heuristics that I would implement next, if they sound reasonable to you guys:

Parallel first : prioritize parallel loops over reduction loops
Dense outer : prioritize the most dense loops first
Sparse outer : the opposite, potentially useful in some cases?

There is another that I am considering, stride/memory aware, which would prioritize loops with better stride patterns (like sequential or linear). Not sure how well this carries over to Sparse Tensor though. Are there any ideas/heuristics that I should definitely try to implement?

As we discussed, I will try to incrementally add heuristics. Sorry for the delay on my end, and thank you so much for the feedback!

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

4 Files Affected:

(modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h (+13-1)
(modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+5)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp (+4-1)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp (+11-1)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 212f7b6f13c26..24f30353e58b5 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -55,6 +55,15 @@ enum class SparseEmitStrategy {
   kDebugInterface, // generate only place-holder for sparse iteration
 };
 
+namespace sparse_tensor {
+
+/// Selects between different loop ordering strategies for sparse tensor 
+enum class LoopOrderingStrategy : unsigned {
+  kDefault, ///< Default strategy (current behavior)
+};
+
+} // namespace sparse_tensor
+
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
@@ -72,10 +81,13 @@ std::unique_ptr<Pass> createSparseAssembler(bool directOut);
 //===----------------------------------------------------------------------===//
 
 void populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                  ReinterpretMapScope scope);
+                                  ReinterpretMapScope scope,
+                                  sparse_tensor::LoopOrderingStrategy strategy = sparse_tensor::LoopOrderingStrategy::kDefault);
 
 std::unique_ptr<Pass> createSparseReinterpretMapPass();
 std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope);
+std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope, 
+                                                     sparse_tensor::LoopOrderingStrategy strategy);
 
 //===----------------------------------------------------------------------===//
 // The PreSparsificationRewriting pass.
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 2513e106f5b06..1b0c7f7583827 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -81,6 +81,11 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> {
          clEnumValN(mlir::ReinterpretMapScope::kExceptGeneric,
                     "except-generic",
                     "Run on operations expect linalg.generic (e.g., foreach)"))}]>,
+    Option<"loopOrderingStrategy", "loop-ordering-strategy", "mlir::sparse_tensor::LoopOrderingStrategy",
+       "mlir::sparse_tensor::LoopOrderingStrategy::kDefault",
+       "Set the loop ordering strategy for sparse tensor dialect", [{llvm::cl::values(
+         clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDefault, "default",
+                    "Default strategy (current behavior)"))}]>,
   ];
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index df9b6cf040efa..06c58d2b13cf2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -788,7 +788,10 @@ struct ForeachOpDemapper
 } // namespace
 
 void mlir::populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                        ReinterpretMapScope scope) {
+                                        ReinterpretMapScope scope,
+                                        sparse_tensor::LoopOrderingStrategy strategy) {
+  (void)strategy; // Suppress unused parameter warning
+  
   if (scope == ReinterpretMapScope::kAll ||
       scope == ReinterpretMapScope::kGenericOnly) {
     patterns.add<GenericOpReinterpretMap, GenericOpScheduler>(
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index 153b9b170e5d3..aa31927e0602c 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -67,12 +67,13 @@ struct SparseReinterpretMap
   SparseReinterpretMap(const SparseReinterpretMap &pass) = default;
   SparseReinterpretMap(const SparseReinterpretMapOptions &options) {
     scope = options.scope;
+    loopOrderingStrategy = options.loopOrderingStrategy;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    populateSparseReinterpretMap(patterns, scope);
+    populateSparseReinterpretMap(patterns, scope, loopOrderingStrategy);
     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
   }
 };
@@ -438,6 +439,15 @@ mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope) {
   return std::make_unique<SparseReinterpretMap>(options);
 }
 
+std::unique_ptr<Pass>
+mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope, 
+                                     sparse_tensor::LoopOrderingStrategy strategy) {
+  SparseReinterpretMapOptions options;
+  options.scope = scope;
+  options.loopOrderingStrategy = strategy;
+  return std::make_unique<SparseReinterpretMap>(options);
+}
+
 std::unique_ptr<Pass> mlir::createPreSparsificationRewritePass() {
   return std::make_unique<PreSparsificationRewritePass>();
 }

llvmbot · 2025-08-21T01:18:28Z

@llvm/pr-subscribers-mlir

Author: Govind Malasani (gmalasan)

Changes

As discussed before, this PR adds the basic infrastructure/boiler plate for loop ordering strategies to be implemented.

If this looks ok, I wanted to also mention some of the heuristics that I would implement next, if they sound reasonable to you guys:

Parallel first : prioritize parallel loops over reduction loops
Dense outer : prioritize the most dense loops first
Sparse outer : the opposite, potentially useful in some cases?

There is another that I am considering, stride/memory aware, which would prioritize loops with better stride patterns (like sequential or linear). Not sure how well this carries over to Sparse Tensor though. Are there any ideas/heuristics that I should definitely try to implement?

As we discussed, I will try to incrementally add heuristics. Sorry for the delay on my end, and thank you so much for the feedback!

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

4 Files Affected:

(modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h (+13-1)
(modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+5)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp (+4-1)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp (+11-1)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 212f7b6f13c26..24f30353e58b5 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -55,6 +55,15 @@ enum class SparseEmitStrategy {
   kDebugInterface, // generate only place-holder for sparse iteration
 };
 
+namespace sparse_tensor {
+
+/// Selects between different loop ordering strategies for sparse tensor 
+enum class LoopOrderingStrategy : unsigned {
+  kDefault, ///< Default strategy (current behavior)
+};
+
+} // namespace sparse_tensor
+
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
@@ -72,10 +81,13 @@ std::unique_ptr<Pass> createSparseAssembler(bool directOut);
 //===----------------------------------------------------------------------===//
 
 void populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                  ReinterpretMapScope scope);
+                                  ReinterpretMapScope scope,
+                                  sparse_tensor::LoopOrderingStrategy strategy = sparse_tensor::LoopOrderingStrategy::kDefault);
 
 std::unique_ptr<Pass> createSparseReinterpretMapPass();
 std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope);
+std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope, 
+                                                     sparse_tensor::LoopOrderingStrategy strategy);
 
 //===----------------------------------------------------------------------===//
 // The PreSparsificationRewriting pass.
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 2513e106f5b06..1b0c7f7583827 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -81,6 +81,11 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> {
          clEnumValN(mlir::ReinterpretMapScope::kExceptGeneric,
                     "except-generic",
                     "Run on operations expect linalg.generic (e.g., foreach)"))}]>,
+    Option<"loopOrderingStrategy", "loop-ordering-strategy", "mlir::sparse_tensor::LoopOrderingStrategy",
+       "mlir::sparse_tensor::LoopOrderingStrategy::kDefault",
+       "Set the loop ordering strategy for sparse tensor dialect", [{llvm::cl::values(
+         clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDefault, "default",
+                    "Default strategy (current behavior)"))}]>,
   ];
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index df9b6cf040efa..06c58d2b13cf2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -788,7 +788,10 @@ struct ForeachOpDemapper
 } // namespace
 
 void mlir::populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                        ReinterpretMapScope scope) {
+                                        ReinterpretMapScope scope,
+                                        sparse_tensor::LoopOrderingStrategy strategy) {
+  (void)strategy; // Suppress unused parameter warning
+  
   if (scope == ReinterpretMapScope::kAll ||
       scope == ReinterpretMapScope::kGenericOnly) {
     patterns.add<GenericOpReinterpretMap, GenericOpScheduler>(
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index 153b9b170e5d3..aa31927e0602c 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -67,12 +67,13 @@ struct SparseReinterpretMap
   SparseReinterpretMap(const SparseReinterpretMap &pass) = default;
   SparseReinterpretMap(const SparseReinterpretMapOptions &options) {
     scope = options.scope;
+    loopOrderingStrategy = options.loopOrderingStrategy;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    populateSparseReinterpretMap(patterns, scope);
+    populateSparseReinterpretMap(patterns, scope, loopOrderingStrategy);
     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
   }
 };
@@ -438,6 +439,15 @@ mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope) {
   return std::make_unique<SparseReinterpretMap>(options);
 }
 
+std::unique_ptr<Pass>
+mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope, 
+                                     sparse_tensor::LoopOrderingStrategy strategy) {
+  SparseReinterpretMapOptions options;
+  options.scope = scope;
+  options.loopOrderingStrategy = strategy;
+  return std::make_unique<SparseReinterpretMap>(options);
+}
+
 std::unique_ptr<Pass> mlir::createPreSparsificationRewritePass() {
   return std::make_unique<PreSparsificationRewritePass>();
 }

…egy fully available in the topoSort method in IterationGraphSorter

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

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

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

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h

gmalasan · 2025-08-26T20:30:22Z

Hey @aartbik, thank you so much for the feedback! I believe I've addressed the issues you pointed out, I'd love to know if there's anything else I should fix.

Otherwise, I'd like to ask if what I mentioned earlier for my next steps make sense, implementing a parallel first, dense outer, and a sparse outer heuristic.

I would love to hear your thoughts about these strategies, or honestly if there are any heuristics you think are worth implementing. Thanks again!

aartbik

two last nits.

as for next steps, once this PR goes in, please start incrementally adding new strategies, but make sure to provide evidence of usefulness of each, and avoid hard coding values that are very specific for the experiments you conducted (or at least document that properly)

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

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h

aartbik · 2025-09-08T21:55:25Z

@gmalasan did you see my comments? of course, take your time, but given the silence after the review I was wondering if you got the notice?

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging other changes with comment style fixes

gmalasan · 2025-09-08T22:34:55Z

Sorry for the delay, and thanks for the reminder, I just made the fixes!

I also wanted to ask what evidence I should provide for the usefulness of each heuristic. What kinds of experiments are good to do? Benchmarking with various loop orderings? And is there any infrastructure that already exists for benchmarking this kind of thing? Or somewhere I can check for more help?

Thanks again for all the help!

aartbik · 2025-09-09T16:36:41Z

I also wanted to ask what evidence I should provide for the usefulness of each heuristic. What kinds of experiments are good to do? Benchmarking with various loop orderings? And is there any infrastructure that already exists for benchmarking this kind of thing? Or somewhere I can check for more help?

There is no good single answer for this. MLIR is a compiler infrastructure, and not a specific compiler for a specific target architecture. So in that sense, we should not fine-tune too much for specific CPUs or GPUs. Having said that, in your previous PR you had a lot of hard-coded constants that, I assume, were the result of tuning for a particular target. Let's avoid that but instead document the reasoning for heuristics, like "if there are more reads than writes it is beneficial to.... because...." but not "if (s > 0.2313) heuristic1". It is okay to benchmark specific examples, and validate and even document the outcome. But eventually the code should reflect some reasoning others can follow. Does that make sense?

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

gmalasan · 2025-09-22T01:48:28Z

Thanks for the advice. Sorry for the delay in implementing heuristics on my end. I tried implementing a generic heuristic for preferring dense loops outer and sparse loops/singleton loops inner. Does this heuristic make sense? And as always is there anything I should fix or look out for? I avoided completely using any magic constants.

aartbik · 2025-09-22T16:18:15Z

Wait, are you continuing with new development in this PR? The idea was that we reviewed and approved the flag change, so that can go into main. After that, you send out new PRs (with much smaller deltas) with the new heuristics that we discuss in turn then.

If you just build on top of this PR we again get very huge differences and I have to constantly filter what is new and old?

gmalasan · 2025-09-22T18:26:12Z

Sorry about that, I reset this branch to the last state that you approved and I'm opening a new PR for this heuristic. Thank you for letting me know the process, it definitely makes sense not to do so much in one PR.

aartbik

Great, yes, please submit this and then make a new PR for the first new heuristic, which we will discuss there!

gmalasan · 2025-09-23T18:33:40Z

Hi @aartbik, sorry, I'm a little confused on what you meant by submitting this. I'm seeing that there are two workflows awaiting approval from a maintainer, but I'm not sure if there's anything else I can do. Or is there a way for me to merge the changes into main (I'm assuming not)?

github-actions · 2025-09-23T23:33:40Z

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff origin/main HEAD --extensions cpp,h -- mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h

⚠️
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing origin/main to the base branch/commit you want to compare against.
⚠️

View the diff from clang-format here.

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 9a966a4dd..af64370a6 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -59,7 +59,8 @@ namespace sparse_tensor {
 
 /// Defines a strategy for loop ordering during sparse code generation.
 enum class LoopOrderingStrategy : unsigned {
-  kDefault, ///< Default strategy (eagerly selects last loop in topological sort).
+  kDefault, ///< Default strategy (eagerly selects last loop in topological
+            ///< sort).
 };
 
 } // namespace sparse_tensor
@@ -80,14 +81,16 @@ std::unique_ptr<Pass> createSparseAssembler(bool directOut);
 // The SparseReinterpretMap pass.
 //===----------------------------------------------------------------------===//
 
-void populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                  ReinterpretMapScope scope,
-                                  sparse_tensor::LoopOrderingStrategy strategy = sparse_tensor::LoopOrderingStrategy::kDefault);
+void populateSparseReinterpretMap(
+    RewritePatternSet &patterns, ReinterpretMapScope scope,
+    sparse_tensor::LoopOrderingStrategy strategy =
+        sparse_tensor::LoopOrderingStrategy::kDefault);
 
 std::unique_ptr<Pass> createSparseReinterpretMapPass();
 std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope);
-std::unique_ptr<Pass> createSparseReinterpretMapPass(ReinterpretMapScope scope, 
-                                                     sparse_tensor::LoopOrderingStrategy strategy);
+std::unique_ptr<Pass>
+createSparseReinterpretMapPass(ReinterpretMapScope scope,
+                               sparse_tensor::LoopOrderingStrategy strategy);
 
 //===----------------------------------------------------------------------===//
 // The PreSparsificationRewriting pass.
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index e068eee05..af2dc8b4d 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -407,9 +407,10 @@ public:
 };
 
 struct GenericOpScheduler : public OpRewritePattern<linalg::GenericOp> {
-  GenericOpScheduler(MLIRContext *context, sparse_tensor::LoopOrderingStrategy strategy)
+  GenericOpScheduler(MLIRContext *context,
+                     sparse_tensor::LoopOrderingStrategy strategy)
       : OpRewritePattern<linalg::GenericOp>(context), strategy(strategy) {}
-  
+
   LogicalResult matchAndRewrite(linalg::GenericOp linalgOp,
                                 PatternRewriter &rewriter) const override {
     if (linalgOp.getNumDpsInits() != 1 || !linalgOp.hasPureTensorSemantics() ||
@@ -792,9 +793,9 @@ struct ForeachOpDemapper
 
 } // namespace
 
-void mlir::populateSparseReinterpretMap(RewritePatternSet &patterns,
-                                        ReinterpretMapScope scope,
-                                        sparse_tensor::LoopOrderingStrategy strategy) {
+void mlir::populateSparseReinterpretMap(
+    RewritePatternSet &patterns, ReinterpretMapScope scope,
+    sparse_tensor::LoopOrderingStrategy strategy) {
   if (scope == ReinterpretMapScope::kAll ||
       scope == ReinterpretMapScope::kGenericOnly) {
     patterns.add<GenericOpReinterpretMap>(patterns.getContext());
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index aa31927e0..b660e2215 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -439,9 +439,8 @@ mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope) {
   return std::make_unique<SparseReinterpretMap>(options);
 }
 
-std::unique_ptr<Pass>
-mlir::createSparseReinterpretMapPass(ReinterpretMapScope scope, 
-                                     sparse_tensor::LoopOrderingStrategy strategy) {
+std::unique_ptr<Pass> mlir::createSparseReinterpretMapPass(
+    ReinterpretMapScope scope, sparse_tensor::LoopOrderingStrategy strategy) {
   SparseReinterpretMapOptions options;
   options.scope = scope;
   options.loopOrderingStrategy = strategy;
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 29596e5fe..73e0f3d28 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -100,7 +100,7 @@ AffineMap IterationGraphSorter::topoSort() {
     // We always prefer a parallel loop over a reduction loop because putting
     // a reduction loop early might make the loop sequence inadmissible.
     auto &it = !parIt.empty() ? parIt : redIt;
-    
+
     // Select loop based on strategy.
     unsigned src;
     switch (strategy) {
@@ -108,7 +108,7 @@ AffineMap IterationGraphSorter::topoSort() {
       src = it.back();
       break;
     }
-    
+
     loopOrder.push_back(src);
     it.pop_back();
     // Update in-degree, and push 0-degree node into worklist.
@@ -130,9 +130,8 @@ AffineMap IterationGraphSorter::topoSort() {
   return AffineMap();
 }
 
-IterationGraphSorter
-IterationGraphSorter::fromGenericOp(linalg::GenericOp genericOp, 
-                                   sparse_tensor::LoopOrderingStrategy strategy) {
+IterationGraphSorter IterationGraphSorter::fromGenericOp(
+    linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy) {
   // Must be a demapped sparse kernel.
   assert(!hasAnyNonIdentityOperandsOrResults(genericOp) &&
          hasAnySparseOperandOrResult(genericOp) &&
@@ -157,7 +156,8 @@ IterationGraphSorter::IterationGraphSorter(
     AffineMap loop2OutLvl, SmallVector<utils::IteratorType> &&iterTypes,
     sparse_tensor::LoopOrderingStrategy strategy)
     : ins(std::move(ins)), loop2InsLvl(std::move(loop2InsLvl)), out(out),
-      loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypes)), strategy(strategy) {
+      loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypes)),
+      strategy(strategy) {
   // One map per tensor.
   assert(loop2InsLvl.size() == ins.size());
   // All the affine maps have the same number of dimensions (loops).
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
index 868a978e9..b2a16e938 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
@@ -13,8 +13,8 @@
 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
 
-#include "mlir/IR/AffineMap.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/IR/AffineMap.h"
 
 namespace mlir {
 
@@ -43,9 +43,11 @@ enum class SortMask : unsigned {
 class IterationGraphSorter {
 public:
   /// Factory method that constructs an iteration graph sorter
-  /// for the given linalg.generic operation with a specific loop ordering strategy.
-  static IterationGraphSorter fromGenericOp(linalg::GenericOp genericOp,
-                                           sparse_tensor::LoopOrderingStrategy strategy);
+  /// for the given linalg.generic operation with a specific loop ordering
+  /// strategy.
+  static IterationGraphSorter
+  fromGenericOp(linalg::GenericOp genericOp,
+                sparse_tensor::LoopOrderingStrategy strategy);
 
   /// Returns a permutation that represents the scheduled loop order.
   /// Note that the returned AffineMap could be null if the kernel
@@ -61,7 +63,8 @@ private:
                        SmallVector<AffineMap> &&loop2InsLvl, Value out,
                        AffineMap loop2OutLvl,
                        SmallVector<utils::IteratorType> &&iterTypes,
-                       sparse_tensor::LoopOrderingStrategy strategy = sparse_tensor::LoopOrderingStrategy::kDefault);
+                       sparse_tensor::LoopOrderingStrategy strategy =
+                           sparse_tensor::LoopOrderingStrategy::kDefault);
 
   // Adds all the constraints in the given loop to level map.
   void addConstraints(Value t, AffineMap loop2LvlMap);

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging changes from main with local formatting changes

aartbik · 2025-09-25T15:50:15Z

Please run the formatter to fix the style issues

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging changes in infrastructure commit formatting with upstream changes

gmalasan · 2025-09-27T03:37:29Z

I believe the formatting issues are now fixed.

[MLIR][SparseTensor] Loop ordering strategy infrastructure (flag)

b9ae621

gmalasan requested review from aartbik, PeimingLiu, yinying-lisa-li and matthias-springer as code owners August 21, 2025 01:17

llvmbot added mlir:sparse Sparse compiler in MLIR mlir labels Aug 21, 2025

[MLIR][SparseTensor] Fixed up the rest of the boilerplate code, strat…

97cec4e

…egy fully available in the topoSort method in IterationGraphSorter

aartbik reviewed Aug 26, 2025

View reviewed changes

gmalasan added 3 commits August 26, 2025 16:22

[MLIR][SparseTensor] Fixed PR feedback about style

3f3661a

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

17943dc

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

6f3012f

gmalasan added 2 commits August 30, 2025 15:08

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

e9a0671

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

3a5d7e2

aartbik reviewed Sep 2, 2025

View reviewed changes

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td Outdated Show resolved Hide resolved

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h Outdated Show resolved Hide resolved

gmalasan added 2 commits September 8, 2025 18:31

[MLIR][SparseTensor] Comment style fixes

d0b96d7

Merge branch 'sparse-tensor-loop-ordering-infrastructure' of https://…

ccc784d

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging other changes with comment style fixes

aartbik reviewed Sep 9, 2025

View reviewed changes

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h Outdated Show resolved Hide resolved

gmalasan added 2 commits September 9, 2025 17:28

[MLIR][SparseTensor] Missed comment style fix in Passes.h

fde6039

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

e80fee7

aartbik approved these changes Sep 10, 2025

View reviewed changes

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

79ee468

aartbik self-requested a review September 22, 2025 16:18

gmalasan force-pushed the sparse-tensor-loop-ordering-infrastructure branch from bd92633 to 79ee468 Compare September 22, 2025 18:23

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

cc7af1b

gmalasan mentioned this pull request Sep 22, 2025

[MLIR][SparseTensor] Dense Outer Loop Ordering Strategy #160168

Open

aartbik approved these changes Sep 23, 2025

View reviewed changes

gmalasan added 2 commits September 23, 2025 14:10

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

b6f5c6b

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

3832d9d

gmalasan added 3 commits September 23, 2025 22:46

Apply clang-format to C++ files for loop ordering infrastructure

fbfa3b6

Merge branch 'sparse-tensor-loop-ordering-infrastructure' of https://…

8630cc7

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging changes from main with local formatting changes

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

db6ed6d

gmalasan added 3 commits September 26, 2025 23:35

Fix clang-format issues in SparseTensorPasses.cpp

57c4088

Merge branch 'sparse-tensor-loop-ordering-infrastructure' of https://…

7e416fc

…github.com/gmalasan/llvm-project into sparse-tensor-loop-ordering-infrastructure Merging changes in infrastructure commit formatting with upstream changes

Merge branch 'main' into sparse-tensor-loop-ordering-infrastructure

c108117

[MLIR][SparseTensor] Loop ordering strategy infrastructure (flag) #154656

Are you sure you want to change the base?

[MLIR][SparseTensor] Loop ordering strategy infrastructure (flag) #154656

Conversation

gmalasan commented Aug 21, 2025

Uh oh!

github-actions bot commented Aug 21, 2025

Uh oh!

llvmbot commented Aug 21, 2025

Uh oh!

llvmbot commented Aug 21, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gmalasan commented Aug 26, 2025

Uh oh!

aartbik left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

aartbik commented Sep 8, 2025

Uh oh!

gmalasan commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aartbik commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

gmalasan commented Sep 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aartbik commented Sep 22, 2025

Uh oh!

gmalasan commented Sep 22, 2025

Uh oh!

aartbik left a comment

Choose a reason for hiding this comment

Uh oh!

gmalasan commented Sep 23, 2025

Uh oh!

github-actions bot commented Sep 23, 2025

Uh oh!

aartbik commented Sep 25, 2025

Uh oh!

gmalasan commented Sep 27, 2025

Uh oh!

Uh oh!

gmalasan commented Sep 8, 2025 •

edited

Loading

aartbik commented Sep 9, 2025 •

edited

Loading

gmalasan commented Sep 22, 2025 •

edited

Loading