[MLIR][SparseTensor] Added Sparse Outer Loop Ordering Strategy #172198
+47
−12
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
gmalasan
requested review from
PeimingLiu,
aartbik,
matthias-springer and
yinying-lisa-li
as code owners
Member
|
@llvm/pr-subscribers-mlir Author: Govind Malasani (gmalasan) ChangesThis 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:
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;
|
Member
|
@llvm/pr-subscribers-mlir-sparse Author: Govind Malasani (gmalasan) ChangesThis 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:
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;
|
|
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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).