[mlir] Make it possible to build a DenseResourceElementsAttr from untyped memory. (#66009)

Exposes the existing `get(ShapedType, StringRef, AsmResourceBlob)`
builder publicly (was protected) and adds a CAPI
`mlirUnmanagedDenseBlobResourceElementsAttrGet`.

While such a generic construction interface is a big help when it comes
to interop, it is also necessary for creating resources that don't have
a standard C type (i.e. f16, the f8s, etc).

Previously reviewed/approved as part of https://reviews.llvm.org/D157064
This commit is contained in:
Stella Laurenzo 2023-09-11 14:10:03 -07:00 committed by GitHub
parent cc2013061e
commit 7055df7b4f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 35 additions and 12 deletions

View File

@ -600,6 +600,13 @@ mlirUnmanagedDenseDoubleResourceElementsAttrGet(MlirType shapedType,
intptr_t numElements,
const double *elements);
/// Unlike the typed accessors above, constructs the attribute with a raw
/// data buffer and no type/alignment checking. Use a more strongly typed
/// accessor if possible.
MLIR_CAPI_EXPORTED MlirAttribute mlirUnmanagedDenseBlobResourceElementsAttrGet(
MlirType shapedType, MlirStringRef name, const void *data,
size_t dataLength);
/// Returns the pos-th value (flat contiguous indexing) of a specific type
/// contained by the given dense resource elements attribute.
MLIR_CAPI_EXPORTED bool

View File

@ -466,21 +466,20 @@ def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements", [
let builders = [
AttrBuilderWithInferredContext<(ins
"ShapedType":$type, "DenseResourceElementsHandle":$handle
)>
];
let extraClassDeclaration = [{
protected:
)>,
/// A builder that inserts a new resource into the builtin dialect's blob
/// manager using the provided blob. The handle of the inserted blob is used
/// when building the attribute. The provided `blobName` is used as a hint
/// for the key of the new handle for the `blob` resource, but may be
/// changed if necessary to ensure uniqueness during insertion.
static DenseResourceElementsAttr get(
ShapedType type, StringRef blobName, AsmResourceBlob blob
);
/// This base class builder does no element type specific size or alignment
/// checking. Use the typed subclasses for more safety unless if performing
/// generic operations.
AttrBuilderWithInferredContext<(ins
"ShapedType":$type, "StringRef":$blobName, "AsmResourceBlob":$blob
)>
];
public:
}];
let skipDefaultBuilders = 1;
}

View File

@ -852,6 +852,14 @@ mlirUnmanagedDenseDoubleResourceElementsAttrGet(MlirType shapedType,
return getDenseResource<DenseF64ResourceElementsAttr>(shapedType, name,
numElements, elements);
}
MLIR_CAPI_EXPORTED MlirAttribute mlirUnmanagedDenseBlobResourceElementsAttrGet(
MlirType shapedType, MlirStringRef name, const void *data,
size_t dataLength) {
return wrap(DenseResourceElementsAttr::get(
llvm::cast<ShapedType>(unwrap(shapedType)), unwrap(name),
UnmanagedAsmResourceBlob::allocateInferAlign(
llvm::ArrayRef(static_cast<const char *>(data), dataLength))));
}
template <typename U, typename T>
static T getDenseResourceVal(MlirAttribute attr, intptr_t pos) {

View File

@ -1118,7 +1118,8 @@ int printBuiltinAttributes(MlirContext ctx) {
const uint8_t *uint8RawData =
(const uint8_t *)mlirDenseElementsAttrGetRawData(uint8Elements);
const int8_t *int8RawData = (const int8_t *)mlirDenseElementsAttrGetRawData(int8Elements);
const int8_t *int8RawData =
(const int8_t *)mlirDenseElementsAttrGetRawData(int8Elements);
const uint32_t *uint32RawData =
(const uint32_t *)mlirDenseElementsAttrGetRawData(uint32Elements);
const int32_t *int32RawData =
@ -1127,7 +1128,8 @@ int printBuiltinAttributes(MlirContext ctx) {
(const uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements);
const int64_t *int64RawData =
(const int64_t *)mlirDenseElementsAttrGetRawData(int64Elements);
const float *floatRawData = (const float *)mlirDenseElementsAttrGetRawData(floatElements);
const float *floatRawData =
(const float *)mlirDenseElementsAttrGetRawData(floatElements);
const double *doubleRawData =
(const double *)mlirDenseElementsAttrGetRawData(doubleElements);
const uint16_t *bf16RawData =
@ -1268,6 +1270,10 @@ int printBuiltinAttributes(MlirContext ctx) {
MlirAttribute doublesBlob = mlirUnmanagedDenseDoubleResourceElementsAttrGet(
mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding),
mlirStringRefCreateFromCString("resource_f64"), 2, doubles);
MlirAttribute blobBlob = mlirUnmanagedDenseBlobResourceElementsAttrGet(
mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64), encoding),
mlirStringRefCreateFromCString("resource_i64_blob"), uints64,
sizeof(uints64));
mlirAttributeDump(uint8Blob);
mlirAttributeDump(uint16Blob);
@ -1279,6 +1285,7 @@ int printBuiltinAttributes(MlirContext ctx) {
mlirAttributeDump(int64Blob);
mlirAttributeDump(floatsBlob);
mlirAttributeDump(doublesBlob);
mlirAttributeDump(blobBlob);
// CHECK: dense_resource<resource_ui8> : tensor<1x2xui8>
// CHECK: dense_resource<resource_ui16> : tensor<1x2xui16>
// CHECK: dense_resource<resource_ui32> : tensor<1x2xui32>
@ -1289,6 +1296,7 @@ int printBuiltinAttributes(MlirContext ctx) {
// CHECK: dense_resource<resource_i64> : tensor<1x2xi64>
// CHECK: dense_resource<resource_f32> : tensor<1x2xf32>
// CHECK: dense_resource<resource_f64> : tensor<1x2xf64>
// CHECK: dense_resource<resource_i64_blob> : tensor<1x2xi64>
if (mlirDenseUInt8ResourceElementsAttrGetValue(uint8Blob, 1) != 1 ||
mlirDenseUInt16ResourceElementsAttrGetValue(uint16Blob, 1) != 1 ||
@ -1302,7 +1310,8 @@ int printBuiltinAttributes(MlirContext ctx) {
fabsf(mlirDenseFloatResourceElementsAttrGetValue(floatsBlob, 1) - 1.0f) >
1e-6 ||
fabs(mlirDenseDoubleResourceElementsAttrGetValue(doublesBlob, 1) - 1.0f) >
1e-6)
1e-6 ||
mlirDenseUInt64ResourceElementsAttrGetValue(blobBlob, 1) != 1)
return 23;
MlirLocation loc = mlirLocationUnknownGet(ctx);