[mlir][sparse] lower trivial tensor.cast on identical sparse tensors

Even though tensor.cast is not part of the sparse tensor dialect,
it may be used to cast static dimension sizes to dynamic dimension
sizes for sparse tensors without changing the actual sparse tensor
itself. Those cases should be lowered properly when replacing sparse
tensor types with their opaque pointers. Likewise, no op sparse
conversions are handled by this revision in a similar manner.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D112173
This commit is contained in:
Aart Bik
2021-10-20 12:47:31 -07:00
parent 68cac47c26
commit 1b15160ef3
3 changed files with 51 additions and 11 deletions

View File

@@ -97,11 +97,11 @@ struct SparseTensorConversionPass
RewritePatternSet patterns(ctx);
SparseTensorTypeConverter converter;
ConversionTarget target(*ctx);
target.addIllegalOp<ConvertOp, NewOp, ToIndicesOp, ToPointersOp, ToTensorOp,
ToValuesOp>();
// All dynamic rules below accept new function, call, return, and dimop
// operations as legal output of the rewriting provided that all sparse
// tensor types have been fully rewritten.
// Everything in the sparse dialect must go!
target.addIllegalDialect<SparseTensorDialect>();
// All dynamic rules below accept new function, call, return, and tensor
// dim and cast operations as legal output of the rewriting provided that
// all sparse tensor types have been fully rewritten.
target.addDynamicallyLegalOp<FuncOp>(
[&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
target.addDynamicallyLegalOp<CallOp>([&](CallOp op) {
@@ -112,10 +112,13 @@ struct SparseTensorConversionPass
target.addDynamicallyLegalOp<tensor::DimOp>([&](tensor::DimOp op) {
return converter.isLegal(op.getOperandTypes());
});
target.addDynamicallyLegalOp<tensor::CastOp>([&](tensor::CastOp op) {
return converter.isLegal(op.getOperand().getType());
});
// The following operations and dialects may be introduced by the
// rewriting rules, and are therefore marked as legal.
target.addLegalOp<arith::CmpFOp, arith::CmpIOp, arith::ConstantOp,
arith::IndexCastOp, tensor::CastOp, tensor::ExtractOp>();
arith::IndexCastOp, tensor::ExtractOp>();
target.addLegalDialect<LLVM::LLVMDialect, memref::MemRefDialect,
scf::SCFDialect>();
// Populate with rules and apply rewriting rules.