caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][sparse] unifies sparse_tensor.sort_coo/sort into one operation. (#66722)

Browse Source 
The use cases of the two operations are largely overlapped, let's
simplify it and only use one of them.
This commit is contained in:
Peiming Liu
2023-09-19 17:02:32 -07:00
committed by GitHub
parent 74338bfe0c
commit bfa3bc4378
14 changed files with 267 additions and 678 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
View File

@@ -1206,29 +1206,23 @@ private:
// Retrieve the values-array.
Value y = genToValues(rewriter, loc, src);
const auto encSrc = srcTp.getEncoding();
// Sort the COO tensor so that its elements are ordered via increasing
// coordinates for the storage ordering of the dst tensor. Use SortCoo
// if the COO tensor has the same ordering as the dst tensor.
if (dimRank > 1 && srcTp.hasSameDimToLvl(dstTp)) {
Value xs = genToCoordinatesBuffer(rewriter, loc, src);
rewriter.create<SortCooOp>(
loc, nnz, xs, ValueRange{y}, rewriter.getIndexAttr(dimRank),
rewriter.getIndexAttr(0), SparseTensorSortKind::HybridQuickSort);
} else {
// Gather the coordinates-arrays in the dst tensor storage order.
SmallVector<Value> xs(dstLvlRank);
const Level srcLvlRank = srcTp.getLvlRank();
for (Level srcLvl = 0; srcLvl < srcLvlRank; srcLvl++) {
// FIXME: `toOrigDim` is deprecated
Dimension dim = toOrigDim(encSrc, srcLvl);
// FIXME: `toStoredDim` is deprecated
Level dstLvl = toStoredDim(encDst, dim);
xs[dstLvl] =
genToCoordinates(rewriter, loc, src, srcLvl, /*cooStart=*/0);
}
rewriter.create<SortOp>(loc, nnz, xs, ValueRange{y},
SparseTensorSortKind::HybridQuickSort);
// Builds the dstLvl -> srcLvl permutation maps.
SmallVector<AffineExpr> es(dstLvlRank);
const Level srcLvlRank = srcTp.getLvlRank();
for (Level srcLvl = 0; srcLvl < srcLvlRank; srcLvl++) {
// FIXME: `toOrigDim` is deprecated
Dimension dim = toOrigDim(encSrc, srcLvl);
// FIXME: `toStoredDim` is deprecated
Level dstLvl = toStoredDim(encDst, dim);
es[dstLvl] = rewriter.getAffineDimExpr(srcLvl);
}
auto xPerm = AffineMap::get(dstLvlRank, 0, es, rewriter.getContext());
assert(xPerm.isPermutation()); // must be a permutation.
Value xs = genToCoordinatesBuffer(rewriter, loc, src);
rewriter.create<SortCooOp>(loc, nnz, xs, ValueRange{y}, xPerm,
rewriter.getIndexAttr(0),
SparseTensorSortKind::HybridQuickSort);
}
// For each element in the COO tensor, insert the element to the dst tensor.