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

[mlir][sparse] Cleaning up some usage of SparseTensorType

Browse Source 
This is a followup to D147192.

Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D147196
This commit is contained in:
wren romano
2023-03-29 18:36:24 -07:00
parent 498aa534f4
commit f2696e469a
1 changed files with 12 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -356,16 +356,10 @@ public:
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
Value srcTensor = op.getSrc();
auto srcTp = getRankedTensorType(srcTensor);
auto dstTp = getRankedTensorType(op.getResult());
SparseTensorType srcStt(srcTp);
SparseTensorType dstStt(dstTp);
const auto encSrc = srcStt.getEncoding();
if (!srcStt.hasEncoding() || !dstStt.hasEncoding()) {
const auto srcTp = getSparseTensorType(srcTensor);
const auto dstTp = getSparseTensorType(op.getResult());
if (!srcTp.hasEncoding() || !dstTp.hasEncoding())
return failure();
}
// Generate code to represent the static dimension constants or compute
// the dynamic dimension values.
@@ -373,11 +367,11 @@ public:
sizesForTensor(rewriter, srcSizes, loc, srcTp, srcTensor);
SmallVector<Value> dstSizes;
SmallVector<Value> dstDynSizes;
if (dstTp.hasStaticShape()) {
for (auto d : dstTp.getShape())
if (dstTp.hasStaticDimShape()) {
for (Dimension d : dstTp.getDimShape())
dstSizes.push_back(constantIndex(rewriter, loc, d));
} else {
ArrayRef<int64_t> dstShape = dstTp.getShape();
ArrayRef<DynSize> dstShape = dstTp.getDimShape();
genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
op.getReassociationIndices());
for (auto [idx, shape] : llvm::enumerate(dstShape)) {
@@ -389,8 +383,8 @@ public:
// Only need a unordered COO buffer if input and output are not sorted
// in the same way.
Type bufferTp =
srcStt.isAllOrdered() && srcStt.isIdentity() && dstStt.isIdentity()
? dstTp
srcTp.isAllOrdered() && srcTp.isIdentity() && dstTp.isIdentity()
? dstTp.getRankedTensorType()
: getUnorderedCOOFromType(dstTp);
Value buffer =
@@ -406,11 +400,12 @@ public:
// followed by an optional
// %t = sparse_tensor.cast %tmp
// depending on whether the input/output are sorted in the same way.
const auto encSrc = srcTp.getEncoding();
ForeachOp foreachOp = rewriter.create<ForeachOp>(
loc, srcTensor, buffer,
[&](OpBuilder &builder, Location loc, ValueRange srcLcvs, Value v,
ValueRange reduc) {
const Dimension dimRank = srcTp.getRank();
const Dimension dimRank = srcTp.getDimRank();
SmallVector<Value> srcDcvs;
srcDcvs.reserve(dimRank);
for (Dimension d = 0; d < dimRank; d++) {
@@ -427,7 +422,8 @@ public:
Value t = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
if (bufferTp != dstTp) {
Value converted = rewriter.create<ConvertOp>(loc, dstTp, t).getResult();
auto dstRTT = dstTp.getRankedTensorType();
Value converted = rewriter.create<ConvertOp>(loc, dstRTT, t).getResult();
rewriter.create<DeallocTensorOp>(loc, t);
t = converted;
}