[mlir][sparse] comment cleanup in iteration graph sorter #75508
Merged
Conversation
This file contains 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
|
@llvm/pr-subscribers-mlir-sparse Author: Aart Bik (aartbik) ChangesFull diff: https://github.com/llvm/llvm-project/pull/75508.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 588400f3dbaf09..66f96ba08c0ed2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -29,22 +29,23 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
explicit AffineDimFinder(ArrayRef<utils::IteratorType> itTypes)
: iterTypes(itTypes) {}
- // Override method from AffineExprVisitor.
+ /// Overrides the visit method from AffineExprVisitor.
void visitDimExpr(AffineDimExpr expr) {
if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()])
pickedDim = expr;
}
- /// Set the desired iterator type that we want to pick.
+ /// Sets the desired iterator type that we want to pick.
void setPickedIterType(utils::IteratorType iterType) {
pickIterType = iterType;
}
- /// Get the desired AffineDimExpr.
+ /// Gets the desired AffineDimExpr.
AffineDimExpr getDimExpr() const {
return llvm::cast<AffineDimExpr>(pickedDim);
}
+ /// Walks the graph in post order to find dim expr.
void walkPostOrder(AffineExpr expr) {
pickedDim = nullptr;
AffineExprVisitor<AffineDimFinder>::walkPostOrder(expr);
@@ -55,11 +56,11 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
AffineExpr pickedDim;
/// The iterator type that we want.
utils::IteratorType pickIterType;
- /// The mapping between dim=>iterator type.
+ /// The mapping between levels and iterator types.
ArrayRef<utils::IteratorType> iterTypes;
};
-// Flattens an affine expression into a list of AffineDimExprs.
+/// Flattens an affine expression into a list of AffineDimExprs.
struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
// Overrides method from AffineExprVisitor.
void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
@@ -97,8 +98,8 @@ AffineMap IterationGraphSorter::topoSort() {
SmallVector<unsigned> loopOrder;
while (!redIt.empty() || !parIt.empty()) {
- // We always prefer parallel loop over reduction loop because putting
- // reduction loop early might make the loop sequence inadmissible.
+ // 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;
auto src = it.back();
loopOrder.push_back(src);
@@ -114,6 +115,7 @@ AffineMap IterationGraphSorter::topoSort() {
}
}
+ // Return the topological sort on success.
if (loopOrder.size() == numLoops)
return AffineMap::getPermutationMap(loopOrder, out.getContext());
@@ -164,13 +166,14 @@ IterationGraphSorter::IterationGraphSorter(
}
AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
- // Reset the interation graph.
+ // Reset the adjacency matrix that represents the iteration graph.
for (auto &row : itGraph)
std::fill(row.begin(), row.end(), false);
- // Reset cached in-degree.
+ // Reset in-degree.
std::fill(inDegree.begin(), inDegree.end(), 0);
+ // Add the constraints for the loop to level map.
for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
// Get map and encoding.
const auto enc = getSparseTensorEncoding(in.getType());
@@ -180,11 +183,12 @@ AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
addConstraints(in, map);
}
- // Get map and encoding.
+ // Add the constraints for the output map.
const auto enc = getSparseTensorEncoding(out.getType());
if ((enc || includesDenseOutput(mask)) && out != ignored)
addConstraints(out, loop2OutLvl);
+ // Return the topological sort (empty for cyclic).
return topoSort();
}
@@ -196,6 +200,7 @@ void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
}
};
+ // Set up a reduction finder.
AffineDimFinder finder(iterTypes);
finder.setPickedIterType(utils::IteratorType::reduction);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
index be94bb5dffde63..a6abe9eb76c476 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
@@ -79,7 +79,7 @@ class IterationGraphSorter {
// Loop itation types;
SmallVector<utils::IteratorType> iterTypes;
- // Adjacent matrix that represents the iteration graph.
+ // Adjacency matrix that represents the iteration graph.
std::vector<std::vector<bool>> itGraph;
// InDegree used for topo sort.
|
|
@llvm/pr-subscribers-mlir Author: Aart Bik (aartbik) ChangesFull diff: https://github.com/llvm/llvm-project/pull/75508.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 588400f3dbaf09..66f96ba08c0ed2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -29,22 +29,23 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
explicit AffineDimFinder(ArrayRef<utils::IteratorType> itTypes)
: iterTypes(itTypes) {}
- // Override method from AffineExprVisitor.
+ /// Overrides the visit method from AffineExprVisitor.
void visitDimExpr(AffineDimExpr expr) {
if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()])
pickedDim = expr;
}
- /// Set the desired iterator type that we want to pick.
+ /// Sets the desired iterator type that we want to pick.
void setPickedIterType(utils::IteratorType iterType) {
pickIterType = iterType;
}
- /// Get the desired AffineDimExpr.
+ /// Gets the desired AffineDimExpr.
AffineDimExpr getDimExpr() const {
return llvm::cast<AffineDimExpr>(pickedDim);
}
+ /// Walks the graph in post order to find dim expr.
void walkPostOrder(AffineExpr expr) {
pickedDim = nullptr;
AffineExprVisitor<AffineDimFinder>::walkPostOrder(expr);
@@ -55,11 +56,11 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
AffineExpr pickedDim;
/// The iterator type that we want.
utils::IteratorType pickIterType;
- /// The mapping between dim=>iterator type.
+ /// The mapping between levels and iterator types.
ArrayRef<utils::IteratorType> iterTypes;
};
-// Flattens an affine expression into a list of AffineDimExprs.
+/// Flattens an affine expression into a list of AffineDimExprs.
struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
// Overrides method from AffineExprVisitor.
void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
@@ -97,8 +98,8 @@ AffineMap IterationGraphSorter::topoSort() {
SmallVector<unsigned> loopOrder;
while (!redIt.empty() || !parIt.empty()) {
- // We always prefer parallel loop over reduction loop because putting
- // reduction loop early might make the loop sequence inadmissible.
+ // 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;
auto src = it.back();
loopOrder.push_back(src);
@@ -114,6 +115,7 @@ AffineMap IterationGraphSorter::topoSort() {
}
}
+ // Return the topological sort on success.
if (loopOrder.size() == numLoops)
return AffineMap::getPermutationMap(loopOrder, out.getContext());
@@ -164,13 +166,14 @@ IterationGraphSorter::IterationGraphSorter(
}
AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
- // Reset the interation graph.
+ // Reset the adjacency matrix that represents the iteration graph.
for (auto &row : itGraph)
std::fill(row.begin(), row.end(), false);
- // Reset cached in-degree.
+ // Reset in-degree.
std::fill(inDegree.begin(), inDegree.end(), 0);
+ // Add the constraints for the loop to level map.
for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
// Get map and encoding.
const auto enc = getSparseTensorEncoding(in.getType());
@@ -180,11 +183,12 @@ AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
addConstraints(in, map);
}
- // Get map and encoding.
+ // Add the constraints for the output map.
const auto enc = getSparseTensorEncoding(out.getType());
if ((enc || includesDenseOutput(mask)) && out != ignored)
addConstraints(out, loop2OutLvl);
+ // Return the topological sort (empty for cyclic).
return topoSort();
}
@@ -196,6 +200,7 @@ void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
}
};
+ // Set up a reduction finder.
AffineDimFinder finder(iterTypes);
finder.setPickedIterType(utils::IteratorType::reduction);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
index be94bb5dffde63..a6abe9eb76c476 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
@@ -79,7 +79,7 @@ class IterationGraphSorter {
// Loop itation types;
SmallVector<utils::IteratorType> iterTypes;
- // Adjacent matrix that represents the iteration graph.
+ // Adjacency matrix that represents the iteration graph.
std::vector<std::vector<bool>> itGraph;
// InDegree used for topo sort.
|
PeimingLiu
approved these changes
Dec 14, 2023
yinying-lisa-li
approved these changes
Dec 14, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
No description provided.