[Mlir-commits] [mlir] [mlir] Make it possible to build a DenseResourceElementsAttr from unt… (PR #66009)

[Stella Laurenzo]

https://github.com/stellaraccident updated https://github.com/llvm/llvm-project/pull/66009:

>From ea9c32dd9f5244914f85b0540fd8816694d4fb96 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident at gmail.com>
Date: Mon, 11 Sep 2023 13:45:03 -0700
Subject: [PATCH 2/3] [mlir] Make it possible to build a
 DenseResourceElementsAttr from untyped memory.

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
---
 mlir/include/mlir-c/BuiltinAttributes.h   |  7 +++++++
 mlir/include/mlir/IR/BuiltinAttributes.td | 17 ++++++++---------
 mlir/lib/CAPI/IR/BuiltinAttributes.cpp    |  8 ++++++++
 mlir/test/CAPI/ir.c                       | 15 ++++++++++++---
 4 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir-c/BuiltinAttributes.h b/mlir/include/mlir-c/BuiltinAttributes.h
index 63198192453efb8..93c4ed5692ef26d 100644
--- a/mlir/include/mlir-c/BuiltinAttributes.h
+++ b/mlir/include/mlir-c/BuiltinAttributes.h
@@ -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
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.td b/mlir/include/mlir/IR/BuiltinAttributes.td
index 093303495469c37..390231da662e2d4 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinAttributes.td
@@ -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;
 }
 
diff --git a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
index de221ddbfa7a922..84a958d01d2eb14 100644
--- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
@@ -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) {
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index c3b78fe1762c005..5d78daa296501f4 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -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);