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

[mlir][sparse] Fix rewriting for convert op and concatenate op.

Browse Source 
Fix a problem in convert op rewriting where it used the original index for
ToIndicesOp.

Extend the concatenate op rewriting to handle dense destination and dynamic
shape destination.

Make the concatenate op integration test run on the codegen path.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D138057
This commit is contained in:
bixia1
2022-11-15 13:24:15 -08:00
parent 1447ea059b
commit 555e7835f4
2 changed files with 42 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -15,6 +15,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
@@ -149,9 +150,11 @@ static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
// TODO: Maybe pick the bitwidth based on input/output tensors (probably the
// largest one among them) in the original operation instead of using the
// default value.
unsigned pointerBitWidth = encSrc ? encSrc.getPointerBitWidth() : 0;
unsigned indexBitWidth = encSrc ? encSrc.getIndexBitWidth() : 0;
auto enc = SparseTensorEncodingAttr::get(
ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), AffineMap(),
encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
pointerBitWidth, indexBitWidth);
return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
}
@@ -428,10 +431,24 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
PatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto rtp = op.getType().cast<RankedTensorType>();
// TODO: Build the output shape if needed.
assert(rtp.hasStaticShape());
auto rank = rtp.getRank();
size_t conDim = op.getDimension().getZExtValue();
SmallVector<Value> dynSizes;
if (!rtp.hasStaticShape()) {
ArrayRef<int64_t> rShape = rtp.getShape();
for (const auto &d : llvm::enumerate(rShape)) {
if (d.value() == ShapedType::kDynamicSize) {
Value v =
createOrFoldDimOp(rewriter, loc, op.getOperand(0), d.index());
rewriter.create<tensor::DimOp>(loc, op.getOperand(0), d.index());
for (const auto &opnd : op.getOperands().drop_front()) {
Value t = createOrFoldDimOp(rewriter, loc, opnd, d.index());
v = rewriter.create<arith::AddIOp>(loc, v, t);
}
dynSizes.push_back(v);
}
}
}
// %t = concatenate %s1, %s2, %s3 {dim = 1}
// ==>
// %tmp = bufferization.alloc_tensor : unordered COO
@@ -441,13 +458,11 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
// %t = sparse_tensor.cast %tmp
auto cooTp = getUnorderedCOOFromType(rtp);
auto cooBuffer =
rewriter.create<AllocTensorOp>(loc, cooTp, ValueRange()).getResult();
rewriter.create<AllocTensorOp>(loc, cooTp, dynSizes).getResult();
auto rank = rtp.getRank();
Value offset = constantIndex(rewriter, loc, 0);
ForeachOp foreachOp;
for (Value input : op.getInputs()) {
// Builds the indexing map.
// Build a for op for each input tensor to append new values into the
// output tensor.
foreachOp = rewriter.create<ForeachOp>(
@@ -462,8 +477,16 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
idx = builder.create<arith::AddIOp>(loc, idx, offset);
indices.push_back(idx);
}
auto t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
builder.create<sparse_tensor::YieldOp>(loc, t);
Value cond = genIsNonzero(rewriter, loc, v);
scf::IfOp ifOp = builder.create<scf::IfOp>(
loc, TypeRange(reduc.front().getType()), cond, /*else*/ true);
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
Value t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
rewriter.create<scf::YieldOp>(loc, t);
rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
rewriter.create<scf::YieldOp>(loc, reduc.front());
rewriter.setInsertionPointAfter(ifOp);
rewriter.create<sparse_tensor::YieldOp>(loc, ifOp.getResult(0));
});
// Accumulates the offset. Note that only static-shaped inputs are allowed
// by concatenate op verifier, which saves us from computing the offset
@@ -659,7 +682,7 @@ private:
for (uint64_t i = 0; i < rank; i++) {
uint64_t orgDim = toOrigDim(encSrc, i);
xs[toStoredDim(encDst, orgDim)] = rewriter.create<ToIndicesOp>(
loc, indTp, src, rewriter.getIndexAttr(orgDim));
loc, indTp, src, rewriter.getIndexAttr(i));
}
// Retrieve NNZ.