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

[mlir][sparse] add dense to sparse conversion implementation

Browse Source 
Implements lowering dense to sparse conversion, for static tensor types only.
First step towards general sparse_tensor.convert support.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D107681
This commit is contained in:
Aart Bik
2021-08-06 17:06:46 -07:00
parent ac20e56911
commit 05c7f450df
6 changed files with 335 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -27,6 +27,10 @@ using namespace mlir::sparse_tensor;
namespace {
//===----------------------------------------------------------------------===//
// Helper methods.
//===----------------------------------------------------------------------===//
/// Returns internal type encoding for primary storage. Keep these
/// values consistent with the sparse runtime support library.
static unsigned getPrimaryTypeEncoding(Type tp) {
@@ -105,6 +109,109 @@ static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type result,
return SymbolRefAttr::get(context, name);
}
/// Generates a call into the "swiss army knife" method of the sparse runtime
/// support library for materializing sparse tensors into the computation.
static void genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
SparseTensorEncodingAttr &enc, uint32_t action,
Value ptr) {
Location loc = op->getLoc();
ShapedType resType = op->getResult(0).getType().cast<ShapedType>();
SmallVector<Value, 8> params;
// Sparsity annotations in tensor constant form.
SmallVector<APInt, 4> attrs;
unsigned sz = enc.getDimLevelType().size();
for (unsigned i = 0; i < sz; i++)
attrs.push_back(
APInt(8, getDimLevelTypeEncoding(enc.getDimLevelType()[i])));
params.push_back(getTensor(rewriter, 8, loc, attrs));
// Dimension sizes array of the enveloping *dense* tensor. Useful for either
// verification of external data, or for construction of internal data.
auto shape = resType.getShape();
SmallVector<APInt, 4> sizes;
for (unsigned i = 0; i < sz; i++) {
uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i];
sizes.push_back(APInt(64, s));
}
params.push_back(getTensor(rewriter, 64, loc, sizes));
// Dimension order permutation array. This is the "identity" permutation by
// default, or otherwise the "reverse" permutation of a given ordering, so
// that indices can be mapped quickly to the right position.
SmallVector<APInt, 4> perm(sz);
AffineMap p = enc.getDimOrdering();
if (p) {
assert(p.isPermutation() && p.getNumResults() == sz);
for (unsigned i = 0; i < sz; i++)
perm[p.getDimPosition(i)] = APInt(64, i);
} else {
for (unsigned i = 0; i < sz; i++)
perm[i] = APInt(64, i);
}
params.push_back(getTensor(rewriter, 64, loc, perm));
// Secondary and primary types encoding.
unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
unsigned primary = getPrimaryTypeEncoding(resType.getElementType());
assert(primary);
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secPtr)));
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secInd)));
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(primary)));
// User action and pointer.
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(action)));
params.push_back(ptr);
// Generate the call to create new tensor.
Type ptrType =
LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
StringRef name = "newSparseTensor";
rewriter.replaceOpWithNewOp<CallOp>(
op, ptrType, getFunc(op, name, ptrType, params), params);
}
/// Generates a call that exposes the data pointer as a void pointer.
// TODO: probing the data pointer directly is a bit raw; we should replace
// this with proper memref util calls once they become available.
static bool genPtrCall(ConversionPatternRewriter &rewriter, Operation *op,
Value val, Value &ptr) {
Location loc = op->getLoc();
ShapedType sType = op->getResult(0).getType().cast<ShapedType>();
Type eltType = sType.getElementType();
// Specialize name for the data type. Even though the final buffferized
// version only operates on pointers, different names are required to
// ensure type correctness for all intermediate states.
StringRef name;
if (eltType.isF64())
name = "getPtrF64";
else if (eltType.isF32())
name = "getPtrF32";
else if (eltType.isInteger(64))
name = "getPtrI64";
else if (eltType.isInteger(32))
name = "getPtrI32";
else if (eltType.isInteger(16))
name = "getPtrI16";
else if (eltType.isInteger(8))
name = "getPtrI8";
else
return false;
auto memRefTp = MemRefType::get(sType.getShape(), eltType);
auto unrankedTp = UnrankedMemRefType::get(eltType, 0);
Value c = rewriter.create<memref::BufferCastOp>(loc, memRefTp, val);
Value d = rewriter.create<memref::CastOp>(loc, unrankedTp, c);
Type ptrType =
LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
auto call =
rewriter.create<CallOp>(loc, ptrType, getFunc(op, name, ptrType, d), d);
ptr = call.getResult(0);
return true;
}
//===----------------------------------------------------------------------===//
// Conversion rules.
//===----------------------------------------------------------------------===//
/// Sparse conversion rule for returns.
class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
public:
@@ -141,56 +248,11 @@ class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
LogicalResult
matchAndRewrite(NewOp op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const override {
Location loc = op.getLoc();
Type resType = op.getType();
Type eltType = resType.cast<ShapedType>().getElementType();
MLIRContext *context = op->getContext();
SmallVector<Value, 5> params;
// Sparse encoding.
auto enc = getSparseTensorEncoding(resType);
if (!enc)
return failure();
// User pointer.
params.push_back(operands[0]);
// Sparsity annotations in tensor constant form.
SmallVector<APInt, 4> attrs;
unsigned sz = enc.getDimLevelType().size();
for (unsigned i = 0; i < sz; i++)
attrs.push_back(
APInt(8, getDimLevelTypeEncoding(enc.getDimLevelType()[i])));
params.push_back(getTensor(rewriter, 8, loc, attrs));
// Dimension order permutation array. This is the "identity"
// permutation by default, or otherwise the "reverse" permutation
// of a given ordering, so that indices can be mapped quickly
// to the right position.
SmallVector<APInt, 4> perm(sz);
AffineMap p = enc.getDimOrdering();
if (p) {
assert(p.isPermutation() && p.getNumResults() == sz);
for (unsigned i = 0; i < sz; i++)
perm[p.getDimPosition(i)] = APInt(64, i);
} else {
for (unsigned i = 0; i < sz; i++)
perm[i] = APInt(64, i);
}
params.push_back(getTensor(rewriter, 64, loc, perm));
// Secondary and primary types encoding.
unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
unsigned primary = getPrimaryTypeEncoding(eltType);
if (!primary)
return failure();
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secPtr)));
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secInd)));
params.push_back(
rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(primary)));
// Generate the call to create new tensor.
Type ptrType = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
StringRef name = "newSparseTensor";
rewriter.replaceOpWithNewOp<CallOp>(
op, ptrType, getFunc(op, name, ptrType, params), params);
genNewCall(rewriter, op, enc, 0, operands[0]);
return success();
}
};
@@ -201,8 +263,19 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
LogicalResult
matchAndRewrite(ConvertOp op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const override {
// TODO: implement conversions lowering
return failure();
Type resType = op.getType();
auto encDst = getSparseTensorEncoding(resType);
auto encSrc = getSparseTensorEncoding(op.source().getType());
// TODO: implement sparse => sparse
// and sparse => dense
if (!encDst || encSrc)
return failure();
// This is a dense => sparse conversion.
Value ptr;
if (!genPtrCall(rewriter, op, operands[0], ptr))
return failure();
genNewCall(rewriter, op, encDst, 1, ptr);
return success();
}
};
@@ -325,6 +398,10 @@ public:
} // namespace
//===----------------------------------------------------------------------===//
// Public method for populating conversion rules.
//===----------------------------------------------------------------------===//
/// Populates the given patterns list with conversion rules required for
/// the sparsification of linear algebra operations.
void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,