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

[mlir][sparse] start using size_hint provided in allocation op

Browse Source 
Even though we introduced the size_hint, we never used it.
This is a very first step, using the hint during the codegen path.
Note that we can refine the heuristics. Also, we need to start
adding the hint on all allocation generated for reading tensors,
converting tensors, etc.

Reviewed By: Peiming, bixia

Differential Revision: https://reviews.llvm.org/D143292
This commit is contained in:
Aart Bik
2023-02-03 13:08:46 -08:00
parent be58d484cb
commit e2e6e7a6a3
2 changed files with 43 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -198,8 +198,10 @@ static Value createAllocation(OpBuilder &builder, Location loc,
///
static void createAllocFields(OpBuilder &builder, Location loc, Type type,
ValueRange dynSizes, bool enableInit,
SmallVectorImpl<Value> &fields) {
SmallVectorImpl<Value> &fields, Value sizeHint) {
RankedTensorType rtp = type.cast<RankedTensorType>();
SparseTensorEncodingAttr enc = getSparseTensorEncoding(rtp);
// Build original sizes.
SmallVector<Value> sizes;
auto shape = rtp.getShape();
@@ -211,19 +213,34 @@ static void createAllocFields(OpBuilder &builder, Location loc, Type type,
sizes.push_back(constantIndex(builder, loc, shape[r]));
}
Value heuristic = constantIndex(builder, loc, 16);
Value valHeuristic = heuristic;
SparseTensorEncodingAttr enc = getSparseTensorEncoding(rtp);
// Set up some heuristic sizes. We try to set the initial
// size based on available information. Otherwise we just
// initialize a few elements to start the reallocation chain.
// TODO: refine this
Value ptrHeuristic, idxHeuristic, valHeuristic;
if (enc.isAllDense()) {
Value linear = sizes[0];
for (unsigned r = 1; r < rank; r++) {
linear = builder.create<arith::MulIOp>(loc, linear, sizes[r]);
}
valHeuristic = linear;
} else if (sizeHint) {
if (getCOOStart(enc) == 0) {
ptrHeuristic = constantIndex(builder, loc, 2);
idxHeuristic = builder.create<arith::MulIOp>(
loc, constantIndex(builder, loc, rank), sizeHint); // AOS
} else {
ptrHeuristic = idxHeuristic = constantIndex(builder, loc, 16);
}
valHeuristic = sizeHint;
} else {
ptrHeuristic = idxHeuristic = valHeuristic =
constantIndex(builder, loc, 16);
}
foreachFieldAndTypeInSparseTensor(
rtp,
[&builder, &fields, rtp, loc, heuristic, valHeuristic,
[&builder, &fields, rtp, loc, ptrHeuristic, idxHeuristic, valHeuristic,
enableInit](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
unsigned /*dim*/, DimLevelType /*dlt*/) -> bool {
assert(fields.size() == fIdx);
@@ -235,11 +252,12 @@ static void createAllocFields(OpBuilder &builder, Location loc, Type type,
case SparseTensorFieldKind::PtrMemRef:
case SparseTensorFieldKind::IdxMemRef:
case SparseTensorFieldKind::ValMemRef:
field = createAllocation(builder, loc, fType.cast<MemRefType>(),
fKind == SparseTensorFieldKind::ValMemRef
? valHeuristic
: heuristic,
enableInit);
field = createAllocation(
builder, loc, fType.cast<MemRefType>(),
(fKind == SparseTensorFieldKind::PtrMemRef) ? ptrHeuristic
: (fKind == SparseTensorFieldKind::IdxMemRef) ? idxHeuristic
: valHeuristic,
enableInit);
break;
}
assert(field);
@@ -691,9 +709,10 @@ public:
// Construct allocation for each field.
Location loc = op.getLoc();
Value sizeHint = op.getSizeHint();
SmallVector<Value> fields;
createAllocFields(rewriter, loc, resType, adaptor.getOperands(),
enableBufferInitialization, fields);
enableBufferInitialization, fields, sizeHint);
// Replace operation with resulting memrefs.
rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields));
return success();