intel/llvm
1
0
Fork 0
mirror of https://github.com/intel/llvm.git synced 2026-01-17 23:45:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

[mlir][sparse] allow foreach operation to generate out-of-order loop on non-annotated tensor.

Browse Source 
No need for a temp COO and sort even when converting dense -> CSC, we can instead rotate the loop to yield a ordered coordinates at beginning.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144213
This commit is contained in:
Peiming Liu
2023-02-16 20:24:01 +00:00
parent 2062e90aa5
commit 9e8d9316ce
7 changed files with 107 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View File

@@ -982,17 +982,10 @@ LogicalResult CompressOp::verify() {
void ForeachOp::build(
OpBuilder &builder, OperationState &result, Value tensor,
ValueRange initArgs, AffineMapAttr order,
function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
bodyBuilder) {
build(builder, result, tensor, std::nullopt, bodyBuilder);
}
void ForeachOp::build(
OpBuilder &builder, OperationState &result, Value tensor,
ValueRange initArgs,
function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
bodyBuilder) {
build(builder, result, initArgs.getTypes(), tensor, initArgs);
build(builder, result, initArgs.getTypes(), tensor, initArgs, order);
// Builds foreach body.
if (!bodyBuilder)
return;
@@ -1023,6 +1016,10 @@ LogicalResult ForeachOp::verify() {
const Dimension dimRank = t.getDimRank();
const auto args = getBody()->getArguments();
if (getOrder().has_value() &&
(t.getEncoding() || !getOrder()->isPermutation()))
return emitError("Only support permuted order on non encoded dense tensor");
if (static_cast<size_t>(dimRank) + 1 + getInitArgs().size() != args.size())
return emitError("Unmatched number of arguments in the block");