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

[mlir][sparse] implement sparse_tensor.reorder_coo (#68916)

Browse Source 
As a side effect of the change, it also unifies the convertOp
implementation between lib/codegen path.
This commit is contained in:
Peiming Liu
2023-10-12 13:22:45 -07:00
committed by GitHub
parent 220244b71b
commit f248d0b28d
22 changed files with 265 additions and 1345 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
View File

@@ -680,31 +680,26 @@ public:
};
// TODO: use a new SortCOO operation here instead of reusing convert op.
struct SparseSortCOOConverter : public OpConversionPattern<ConvertOp> {
struct SparseReorderCOOConverter : public OpConversionPattern<ReorderCOOOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(ConvertOp op, ConvertOpAdaptor adaptor,
matchAndRewrite(ReorderCOOOp op, ReorderCOOOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// Direct conversion should have already been lowered.
if (!op.isSortCOOConvert())
return failure();
Location loc = op.getLoc();
MLIRContext *ctx = op.getContext();
SparseTensorType srcStt = getSparseTensorType(op.getSource());
SparseTensorType dstStt = getSparseTensorType(op.getDest());
SparseTensorType srcStt = getSparseTensorType(op.getInputCoo());
SparseTensorType dstStt = getSparseTensorType(op.getResultCoo());
// TODO: This should be verification rules for sort_coo operation.
// Should have been verified.
assert(dstStt.isAllOrdered() && !srcStt.isAllOrdered() &&
isUniqueCOOType(srcStt.getRankedTensorType()) &&
isUniqueCOOType(dstStt.getRankedTensorType()));
assert(dstStt.hasSameDimToLvl(srcStt));
// We don't need a mutable descriptor here as we perform sorting in-place.
auto nnz = genValMemSize(rewriter, op.getLoc(), adaptor.getSource());
auto desc = getDescriptorFromTensorTuple(adaptor.getSource());
auto nnz = genValMemSize(rewriter, op.getLoc(), adaptor.getInputCoo());
auto desc = getDescriptorFromTensorTuple(adaptor.getInputCoo());
auto crd = desc.getAOSMemRef();
auto val = desc.getValMemRef();
@@ -715,12 +710,11 @@ struct SparseSortCOOConverter : public OpConversionPattern<ConvertOp> {
auto id = AffineMap::getMultiDimIdentityMap(srcStt.getLvlRank(), ctx);
rewriter.create<SortOp>(loc, nnz, crd, ValueRange{val}, id,
rewriter.getIndexAttr(0),
SparseTensorSortKind::HybridQuickSort);
rewriter.getIndexAttr(0), op.getAlgorithm());
// Since we do in-place sorting, the destinate tensor will have the same set
// of memrefs as the source tensor.
rewriter.replaceOp(op, adaptor.getSource());
rewriter.replaceOp(op, adaptor.getInputCoo());
return success();
}
};
@@ -1147,9 +1141,6 @@ public:
LogicalResult
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (op.isSortCOOConvert())
return failure();
SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
SparseTensorEncodingAttr encSrc =
getSparseTensorEncoding(op.getSource().getType());
@@ -1603,7 +1594,7 @@ void mlir::populateSparseTensorCodegenPatterns(
SparseCastConverter, SparseExtractSliceConverter,
SparseTensorLoadConverter, SparseExpandConverter,
SparseCompressConverter, SparseInsertConverter,
SparseSortCOOConverter,
SparseReorderCOOConverter,
SparseSliceGetterOpConverter<ToSliceOffsetOp,
StorageSpecifierKind::DimOffset>,
SparseSliceGetterOpConverter<ToSliceStrideOp,