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

[mlir][sparse] Support sparse2sparse collapse for dynamic sizes

Browse Source 
This patch implements sparse2sparse collapse for operands with dynamic shape.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D131599
This commit is contained in:
Anlun Xu
2022-09-27 16:25:20 -07:00
parent a759477222
commit fad84c3dbe
3 changed files with 201 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
View File

@@ -478,20 +478,20 @@ static bool canUseDirectConversion(
static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
ArrayRef<ReassociationIndices> reassociation,
TensorType dstTp, TensorType srcTp, Value dstIdx,
Value srcIdx) {
Value srcIdx, ArrayRef<Value> dstShape,
ArrayRef<Value> srcShape) {
unsigned dstRank = dstTp.getRank();
unsigned srcRank = srcTp.getRank();
unsigned start = 0;
unsigned i = 0;
bool isExpand = srcRank > dstRank;
ArrayRef<int64_t> shape = isExpand ? srcTp.getShape() : dstTp.getShape();
ArrayRef<Value> shape = isExpand ? srcShape : dstShape;
// Iterate over reassociation map.
for (const auto &map : llvm::enumerate(reassociation)) {
// Prepare strides information in dimension slice.
uint64_t linear = 1;
Value linear = constantIndex(rewriter, loc, 1);
for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
assert(!ShapedType::isDynamic(shape[j]));
linear *= shape[j];
linear = rewriter.create<arith::MulIOp>(loc, linear, shape[j]);
}
// Start collapse.
Value idx = constantIndex(rewriter, loc, i++);
@@ -500,22 +500,17 @@ static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
val = rewriter.create<memref::LoadOp>(loc, srcIdx, idx);
// Iterate over dimension slice.
for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
linear /= shape[j];
Value stride = constantIndex(rewriter, loc, linear);
linear = rewriter.create<arith::DivUIOp>(loc, linear, shape[j]);
Value jdx = constantIndex(rewriter, loc, j);
if (isExpand) {
Value old = rewriter.create<memref::LoadOp>(loc, srcIdx, jdx);
Value mul = linear == 1
? old
: rewriter.create<arith::MulIOp>(loc, old, stride);
Value mul = rewriter.create<arith::MulIOp>(loc, old, linear);
val = val ? rewriter.create<arith::AddIOp>(loc, val, mul) : mul;
} else {
Value old = val;
if (linear != 1)
val = rewriter.create<arith::DivUIOp>(loc, val, stride);
val = rewriter.create<arith::DivUIOp>(loc, val, linear);
rewriter.create<memref::StoreOp>(loc, val, dstIdx, jdx);
if (linear != 1)
val = rewriter.create<arith::RemUIOp>(loc, old, stride);
val = rewriter.create<arith::RemUIOp>(loc, old, linear);
}
}
// Finalize expansion.
@@ -527,6 +522,65 @@ static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
assert((isExpand && i == dstRank) || (!isExpand && i == srcRank));
}
/// Helper method to compute the shape of destination tensor of a reshape
/// operator. This is only used when operands have dynamic shape. The shape of
/// the destination is stored into dstShape.
void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
SmallVector<Value, 4> &dstShape,
ArrayRef<Value> srcShape,
ArrayRef<int64_t> staticDstShape,
ArrayRef<ReassociationIndices> reassociation) {
// Collapse shape.
if (reassociation.size() < srcShape.size()) {
unsigned start = 0;
for (const auto &map : llvm::enumerate(reassociation)) {
auto dstDim = constantIndex(rewriter, loc, 1);
for (unsigned i = start; i < start + map.value().size(); i++) {
dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
}
dstShape.push_back(dstDim);
start = start + map.value().size();
}
assert(start == srcShape.size());
return;
}
// Expand shape.
assert(reassociation.size() == srcShape.size());
unsigned start = 0;
// Expand the i-th dimension in srcShape.
for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
auto map = reassociation[i];
auto srcDim = srcShape[i];
// Iterate through dimensions expanded from the i-th dimension.
for (unsigned j = start; j < start + map.size(); j++) {
// There can be only one dynamic sized dimension among dimensions expanded
// from the i-th dimension in srcShape. For example, if srcDim = 8, then
// the expanded shape could be <2x?x2>, but not <2x?x?>.
if (staticDstShape[j] == ShapedType::kDynamicSize) {
// The expanded dimension has dynamic size. We compute the dimension
// by dividing srcDim by the product of the static dimensions.
int64_t product = 1;
for (unsigned k = start; k < start + map.size(); k++) {
if (staticDstShape[k] != ShapedType::kDynamicSize) {
product *= staticDstShape[k];
}
}
// Compute the dynamic dimension size.
Value productVal = constantIndex(rewriter, loc, product);
Value dynamicSize =
rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
dstShape.push_back(dynamicSize);
} else {
// The expanded dimension is statically known.
dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
}
}
start = start + map.size();
}
assert(start == staticDstShape.size());
}
/// Generate code for a general sparse to sparse reshaping operation.
/// Note that unlike dense reshaping (which can be done with a "cheap"
/// change of view), sparse reshaping is currently done with actual
@@ -562,19 +616,23 @@ genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
auto noPerm = SparseTensorEncodingAttr::get(
op->getContext(), encSrc.getDimLevelType(), AffineMap(),
encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
SmallVector<Value, 4> sizes;
SmallVector<Value, 4> srcSizes;
SmallVector<Value, 8> params;
sizesFromPtr(rewriter, sizes, loc, encSrc, srcTp, adaptor.getSrc());
newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, sizes,
sizesFromPtr(rewriter, srcSizes, loc, encSrc, srcTp, adaptor.getSrc());
newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, srcSizes,
adaptor.getSrc());
Value iter = genNewCall(rewriter, loc, params);
// Start a new COO for the destination tensor.
sizes.clear();
SmallVector<Value, 4> dstSizes;
params.clear();
// Fills sizes array using the sizes from destination type.
assert(dstTp.hasStaticShape());
sizesFromType(rewriter, sizes, loc, dstTp);
newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, sizes);
if (dstTp.hasStaticShape()) {
sizesFromType(rewriter, dstSizes, loc, dstTp);
} else {
ArrayRef<int64_t> dstShape = dstTp.getShape();
genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
op.getReassociationIndices());
}
newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, dstSizes);
Value coo = genNewCall(rewriter, loc, params);
Value dstPerm = params[2];
// Construct a while loop over the iterator.
@@ -593,7 +651,7 @@ genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
rewriter.setInsertionPointToStart(after);
translateIndices(loc, rewriter, op.getReassociationIndices(), dstTp, srcTp,
dstIdx, srcIdx);
dstIdx, srcIdx, dstSizes, srcSizes);
genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
rewriter.create<scf::YieldOp>(loc);
// Final call to construct sparse tensor storage and free temporary resources.