Skip to content

Commit df11a2b

Browse files
author
Peiming Liu
committed
[mlir][sparse] admit un-sparsifiable operations if all its operands are loaded from dense input
Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D153998
1 parent a484e02 commit df11a2b

File tree

4 files changed

+270
-79
lines changed

4 files changed

+270
-79
lines changed

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ struct TensorExp final {
7070
/// and kSelect, this holds the original operation with all regions. For
7171
/// kBinaryBranch, this holds the YieldOp for the left or right half
7272
/// to be merged into a nested scf loop.
73+
///
74+
/// Or the actual operation that we can not sparsify but having all dense
75+
/// operands for kDenseOp.
7376
Operation *op;
7477

7578
/// An optional attribute that is required to determine the semantics of the
@@ -157,8 +160,9 @@ enum class TensorExp::Kind {
157160
kShrS, // signed
158161
kShrU, // unsigned
159162
kShlI,
160-
kBinary, // semiring binary op
161-
kReduce, // semiring reduction op
163+
kBinary, // semiring binary op
164+
kReduce, // semiring reduction op
165+
kDenseOp, // special category of operations requiring all dense operands
162166
};
163167

164168
//===----------------------------------------------------------------------===//
@@ -645,7 +649,11 @@ class Merger {
645649
Type inferType(ExprId e, Value src) const;
646650

647651
/// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
648-
std::optional<ExprId> buildTensorExp(linalg::GenericOp op, Value v);
652+
/// The boolean value returned indicates whether the result of the current
653+
/// operation being built depends on any value that is loaded from a sparse
654+
/// tensor.
655+
std::pair<std::optional<ExprId>, bool> buildTensorExp(linalg::GenericOp op,
656+
Value v);
649657

650658
/// Merger data structures.
651659
const TensorId outTensor;

0 commit comments

Comments
 (0)