[Mlir-commits] [mlir] [mlir][sparse] remove LoopOrd type (PR #74540)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Dec 5 16:23:08 PST 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Aart Bik (aartbik)
<details>
<summary>Changes</summary>
Rationale:
We no longer deal with topsort during sparsification, so that LoopId == LoopOrd for all methods. This first revision removes the types. A follow up revision will simplify some other remaining constructs that deal with loop order (e.g. at and ldx).
---
Full diff: https://github.com/llvm/llvm-project/pull/74540.diff
4 Files Affected:
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp (+2-3)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h (+3-6)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h (+7-34)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp (+7-7)
``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
index cc05f1d06e30f..312aefc0936c2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
@@ -206,9 +206,8 @@ void CodegenEnv::updateInsertionChain(Value chain) {
insChain = chain;
}
-// FIXME: clarify what this "rank" is really supposed to mean/be.
-bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, LoopOrd n) const {
- return sparseOut == o && outerParNest == static_cast<LoopOrd>(rank - 1) &&
+bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const {
+ return sparseOut == o && outerParNest == static_cast<LoopId>(rank - 1) &&
outerParNest == n;
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
index 7e825dde27830..27dc407b493dc 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
@@ -118,9 +118,7 @@ class CodegenEnv {
/// It also sets the sparseOut if the output tensor is sparse.
bool isAdmissibleTensorExp(ExprId e);
- /// Returns the induction-variable for the loop identified by the given
- /// `LoopId`. This method handles application of the topological sort
- /// in order to convert the `LoopId` into the corresponding `LoopOrd`.
+ /// Returns the induction-variable for the given loop.
Value getLoopVar(LoopId i) const;
//
@@ -133,8 +131,7 @@ class CodegenEnv {
Value getInsertionChain() const { return insChain; }
void updateInsertionChain(Value chain);
- // FIXME: clarify what this "rank" is really supposed to mean/be.
- bool atExpandLevel(OpOperand *o, unsigned rank, LoopOrd n) const;
+ bool atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const;
void startExpand(Value values, Value filled, Value added, Value count);
bool isExpand() const { return expValues != nullptr; }
void updateExpandCount(Value count);
@@ -180,7 +177,7 @@ class CodegenEnv {
// expansion in the innermost loop nest (`expValues` through `expCount`).
OpOperand *sparseOut;
// The count of outer non-filter loops, as defined by `isAdmissibleTopoOrder`.
- LoopOrd outerParNest;
+ LoopId outerParNest;
Value insChain;
Value expValues;
Value expFilled;
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h b/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
index e3e620b92257a..858751932d9d0 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
@@ -19,31 +19,9 @@
namespace mlir {
namespace sparse_tensor {
-//===----------------------------------------------------------------------===//
-/// The position of a loop in the loop-stack, or the position of a
-/// `LoopId` in a topologically-sorted list of `LoopId`s.
-///
-/// Although this type may have the same cardinality as `LoopId`, it must
-/// not be confused with that type. The `LoopId` type is used by the `Merger`
-/// as a unique identifier for loop-variables, regardless of the ordering
-/// of those loops. Whereas the `LoopOrd` type is used by the `LoopEmitter`
-/// (and `CodegenEnv`) to refer to the actual order in which loops are
-/// generated.
-///
-/// TODO: further explicate the correspondences between these various
-/// types. In particular, since the `$dim` argument to `linalg::IndexOp`
-/// is a De Bruijn index, it seems like that should correspond to `LoopOrd`,
-/// and yet the `Merger` has that correspond with `LoopId` instead.
-/// In addition `LoopEmitter::genAffine` has `AffineDimExpr::position`
-/// correspond to `LoopId`, however it is unclear what the providence
-/// of those `AffineDimExpr` is.
-//
-// TODO: use a struct/class rather than a typedef, so that we can actually
-// typecheck this to avoid mixups in the code.
-using LoopOrd = unsigned;
-
// A compressed <tensor id, level> pair.
using TensorLevel = unsigned;
+
//===----------------------------------------------------------------------===//
// SparseTensorLoopEmiter class, manages sparse tensors and helps to
// generate loop structure to (co)-iterate sparse tensors.
@@ -108,9 +86,7 @@ class LoopEmitter {
/// to the position of that tensor `Value` in the array). Setting
/// `isSparseOut` indicates that the sparse output tensor is empty,
/// so the loop emitter will generate loops over it according to the
- /// level-sizes. The `topSort` array specifies the actual order in
- /// which loops are generated, thus providing a mapping from `LoopOrd`
- /// to `LoopId`.
+ /// level-sizes.
void initialize(ValueRange tensors, StringAttr loopTag = nullptr,
bool hasOutput = false, bool isSparseOut = false,
unsigned numLoops = 0, DependentLvlGetter getter = nullptr);
@@ -193,21 +169,18 @@ class LoopEmitter {
}
/// Fills the out-parameter with the loop induction variables for all
- /// loops in the current loop-stack. The variables are given in the
- /// same order as the loop-stack, hence `ivs` should be indexed into
- /// by `LoopOrd` (not `LoopId`).
+ /// loops in the current loop-stack.
SmallVector<Value> getLoopIVs() const {
return llvm::to_vector(getLoopIVsRange());
}
- /// Gets the current depth of the loop-stack. The result is given
- /// the type `LoopOrd` for the same reason as one-past-the-end iterators.
- LoopOrd getCurrentDepth() const {
+ /// Gets the current depth of the loop-stack.
+ LoopId getCurrentDepth() const {
return llvm::range_size(getLoopIVsRange());
}
- /// Gets loop induction variable for the given `LoopOrd`.
- Value getLoopIV(LoopOrd n) const {
+ /// Gets loop induction variable for the given loop
+ Value getLoopIV(LoopId n) const {
if (n >= getCurrentDepth())
return Value();
auto it = getLoopIVsRange().begin();
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index d171087f56ab1..d03e9615d340e 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -411,7 +411,7 @@ static void genInsertionStore(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
Location loc = op.getLoc();
// Direct insertion in lexicographic coordinate order.
if (!env.isExpand()) {
- const LoopOrd numLoops = op.getRank(t);
+ const LoopId numLoops = op.getRank(t);
// Retrieves the first `numLoop` induction variables.
SmallVector<Value> ivs = llvm::to_vector(
llvm::drop_end(env.emitter().getLoopIVsRange(),
@@ -713,7 +713,7 @@ static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
}
/// Generates an expanded access pattern in innermost dimension.
-static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
+static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId at,
bool atStart) {
linalg::GenericOp op = env.op();
OpOperand *lhs = op.getDpsInitOperand(0);
@@ -740,7 +740,7 @@ static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
r.getResult(3));
} else {
SmallVector<Value> indices;
- for (LoopOrd i = 0; i < at; i++)
+ for (LoopId i = 0; i < at; i++)
indices.push_back(env.emitter().getLoopIV(i));
Value values = env.getExpandValues();
Value filled = env.getExpandFilled();
@@ -815,7 +815,7 @@ static Operation *genCoIteration(CodegenEnv &env, OpBuilder &builder,
/// Generates a for-loop or a while-loop, depending on whether it implements
/// singleton iteration or co-iteration over the given conjunction.
-static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
+static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopId at,
bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
bool tryParallel = shouldTryParallize(env, at, at == 0, tidLvls);
return genCoIteration(env, builder, at, tidLvls, tryParallel, needsUniv);
@@ -943,7 +943,7 @@ static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
/// Starts a loop sequence at given level. Returns true if
/// the universal loop index must be maintained at this level.
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
- LoopOrd idx, LoopId ldx, LatSetId lts) {
+ LoopId idx, LoopId ldx, LatSetId lts) {
assert(!env.getLoopVar(idx));
// Emit invariants at this loop sequence level.
genInvariants(env, builder, exp, ldx, /*atStart=*/true);
@@ -1127,7 +1127,7 @@ static bool translateBitsToTidLvlPairs(
/// Starts a single loop in current sequence.
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
- OpBuilder &builder, LoopOrd at,
+ OpBuilder &builder, LoopId at,
LatPointId li, bool needsUniv) {
// The set of tensors + lvls to generate loops on
SmallVector<TensorLevel> tidLvls;
@@ -1199,7 +1199,7 @@ static void endLoopSeq(CodegenEnv &env, OpBuilder &builder, unsigned exp,
/// to manage the complexity of implementing co-iteration over unions
/// and intersections of sparse iterations spaces.
static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp,
- LoopOrd at) {
+ LoopId at) {
// At each leaf, assign remaining tensor (sub)expression to output tensor.
if (at == env.getLoopNum()) {
Value rhs = genExp(env, rewriter, exp);
``````````
</details>
https://github.com/llvm/llvm-project/pull/74540
More information about the Mlir-commits
mailing list