[MLIR] Add function to create BFloat16 array attribute

[ashay-github]

Summary

This patch adds a new function mlirDenseElementsAttrBFloat16Get(),
which accepts the shaped type, the number of BFloat16 values, and a
pointer to an array of BFloat16 values, each of which is a uint16_t
value.

[stellaraccident]

Excellent, thanks. Do you have commit access to land?

[ashay-github]

@stellaraccident Thanks for approving. I do not have commit access yet. Could you (or someone) commit this for now please? Thanks!

mlir/include/mlir-c/BuiltinAttributes.h

 MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrUInt64Get(
     MlirType shapedType, intptr_t numElements, const uint64_t *elements);
 MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrInt64Get(
     MlirType shapedType, intptr_t numElements, const int64_t *elements);
 MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrFloatGet(
     MlirType shapedType, intptr_t numElements, const float *elements);
 MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrDoubleGet(
     MlirType shapedType, intptr_t numElements, const double *elements);
+MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrBFloat16Get(
+    MlirType shapedType, intptr_t numElements, const uint16_t *elements);
 
 /// Creates a dense elements attribute with the given shaped type from string
 /// elements.
 MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrStringGet(
     MlirType shapedType, intptr_t numElements, MlirStringRef *strs);
 
 /// Creates a dense elements attribute that has the same data as the given dense
 /// elements attribute and a different shaped type. The new type must have the

mlir/lib/CAPI/IR/BuiltinAttributes.cpp

MlirAttribute mlirDenseElementsAttrFloatGet(MlirType shapedType,
                                             const float *elements) {
   return getDenseAttribute(shapedType, numElements, elements);
 }
 MlirAttribute mlirDenseElementsAttrDoubleGet(MlirType shapedType,
                                              intptr_t numElements,
                                              const double *elements) {
   return getDenseAttribute(shapedType, numElements, elements);
 }
+MlirAttribute mlirDenseElementsAttrBFloat16Get(MlirType shapedType,
+                                               intptr_t numElements,
+                                               const uint16_t *elements) {
+  size_t bufferSize = numElements * 2;
+  const void *buffer = static_cast<const void *>(elements);
+  return mlirDenseElementsAttrRawBufferGet(shapedType, bufferSize, buffer);
+}
 
 MlirAttribute mlirDenseElementsAttrStringGet(MlirType shapedType,
                                              intptr_t numElements,
                                              MlirStringRef *strs) {
   SmallVector<StringRef, 8> values;
   values.reserve(numElements);
   for (intptr_t i = 0; i < numElements; ++i)
     values.push_back(unwrap(strs[i]));

mlir/test/CAPI/ir.c

int printBuiltinAttributes(MlirContext ctx) {
   uint16_t uints16[] = {0u, 1u};
   int16_t ints16[] = {0, 1};
   uint32_t uints32[] = {0u, 1u};
   int32_t ints32[] = {0, 1};
   uint64_t uints64[] = {0u, 1u};
   int64_t ints64[] = {0, 1};
   float floats[] = {0.0f, 1.0f};
   double doubles[] = {0.0, 1.0};
+  uint16_t bf16s[] = {0x0, 0x3f80};
   MlirAttribute encoding = mlirAttributeGetNull();
   MlirAttribute boolElements = mlirDenseElementsAttrBoolGet(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 1), encoding),
       2, bools);
   MlirAttribute uint8Elements = mlirDenseElementsAttrUInt8Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 8),
                               encoding),
       2, uints8);
MlirAttribute int64Elements = mlirDenseElementsAttrInt64Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64), encoding),
       2, ints64);
   MlirAttribute floatElements = mlirDenseElementsAttrFloatGet(
       mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding), 2,
       floats);
   MlirAttribute doubleElements = mlirDenseElementsAttrDoubleGet(
       mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding), 2,
       doubles);
+  MlirAttribute bf16Elements = mlirDenseElementsAttrBFloat16Get(
+      mlirRankedTensorTypeGet(2, shape, mlirBF16TypeGet(ctx), encoding), 2,
+      bf16s);
 
   if (!mlirAttributeIsADenseElements(boolElements) ||
       !mlirAttributeIsADenseElements(uint8Elements) ||
       !mlirAttributeIsADenseElements(int8Elements) ||
       !mlirAttributeIsADenseElements(uint32Elements) ||
       !mlirAttributeIsADenseElements(int32Elements) ||
       !mlirAttributeIsADenseElements(uint64Elements) ||
       !mlirAttributeIsADenseElements(int64Elements) ||
       !mlirAttributeIsADenseElements(floatElements) ||
-      !mlirAttributeIsADenseElements(doubleElements))
+      !mlirAttributeIsADenseElements(doubleElements) ||
+      !mlirAttributeIsADenseElements(bf16Elements))
     return 14;
 
   if (mlirDenseElementsAttrGetBoolValue(boolElements, 1) != 1 ||
       mlirDenseElementsAttrGetUInt8Value(uint8Elements, 1) != 1 ||
       mlirDenseElementsAttrGetInt8Value(int8Elements, 1) != 1 ||
       mlirDenseElementsAttrGetUInt16Value(uint16Elements, 1) != 1 ||
       mlirDenseElementsAttrGetInt16Value(int16Elements, 1) != 1 ||
       mlirDenseElementsAttrGetUInt32Value(uint32Elements, 1) != 1 ||
int printBuiltinAttributes(MlirContext ctx) {
   mlirAttributeDump(uint8Elements);
   mlirAttributeDump(int8Elements);
   mlirAttributeDump(uint32Elements);
   mlirAttributeDump(int32Elements);
   mlirAttributeDump(uint64Elements);
   mlirAttributeDump(int64Elements);
   mlirAttributeDump(floatElements);
   mlirAttributeDump(doubleElements);
+  mlirAttributeDump(bf16Elements);
   // CHECK: dense<{{\[}}[false, true]]> : tensor<1x2xi1>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xui8>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xi8>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xui32>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xi32>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xui64>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xi64>
   // CHECK: dense<{{\[}}[0.000000e+00, 1.000000e+00]]> : tensor<1x2xf32>
   // CHECK: dense<{{\[}}[0.000000e+00, 1.000000e+00]]> : tensor<1x2xf64>
+  // CHECK: dense<{{\[}}[0.000000e+00, 1.000000e+00]]> : tensor<1x2xbf16>
 
   MlirAttribute splatBool = mlirDenseElementsAttrBoolSplatGet(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 1), encoding),
       1);
   MlirAttribute splatUInt8 = mlirDenseElementsAttrUInt8SplatGet(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 8),
                               encoding),
       1);
int32_t *int32RawData =
       (int32_t *)mlirDenseElementsAttrGetRawData(int32Elements);
   uint64_t *uint64RawData =
       (uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements);
   int64_t *int64RawData =
       (int64_t *)mlirDenseElementsAttrGetRawData(int64Elements);
   float *floatRawData = (float *)mlirDenseElementsAttrGetRawData(floatElements);
   double *doubleRawData =
       (double *)mlirDenseElementsAttrGetRawData(doubleElements);
+  uint16_t *bf16RawData =
+      (uint16_t *)mlirDenseElementsAttrGetRawData(bf16Elements);
   if (uint8RawData[0] != 0u || uint8RawData[1] != 1u || int8RawData[0] != 0 ||
       int8RawData[1] != 1 || uint32RawData[0] != 0u || uint32RawData[1] != 1u ||
       int32RawData[0] != 0 || int32RawData[1] != 1 || uint64RawData[0] != 0u ||
       uint64RawData[1] != 1u || int64RawData[0] != 0 || int64RawData[1] != 1 ||
       floatRawData[0] != 0.0f || floatRawData[1] != 1.0f ||
-      doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0)
+      doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0 ||
+      bf16RawData[0] != 0 || bf16RawData[1] != 0x3f80)
     return 18;
 
   mlirAttributeDump(splatBool);
   mlirAttributeDump(splatUInt8);
   mlirAttributeDump(splatInt8);
   mlirAttributeDump(splatUInt32);
   mlirAttributeDump(splatInt32);
   mlirAttributeDump(splatUInt64);
   mlirAttributeDump(splatInt64);
   mlirAttributeDump(splatFloat);
   mlirAttributeDump(splatDouble);
   // CHECK: dense<true> : tensor<1x2xi1>
   // CHECK: dense<1> : tensor<1x2xui8>
   // CHECK: dense<1> : tensor<1x2xi8>
   // CHECK: dense<1> : tensor<1x2xui32>
   // CHECK: dense<1> : tensor<1x2xi32>
   // CHECK: dense<1> : tensor<1x2xui64>
   // CHECK: dense<1> : tensor<1x2xi64>
   // CHECK: dense<1.000000e+00> : tensor<1x2xf32>
   // CHECK: dense<1.000000e+00> : tensor<1x2xf64>
 
   mlirAttributeDump(mlirElementsAttrGetValue(floatElements, 2, uints64));
   mlirAttributeDump(mlirElementsAttrGetValue(doubleElements, 2, uints64));
+  mlirAttributeDump(mlirElementsAttrGetValue(bf16Elements, 2, uints64));
   // CHECK: 1.000000e+00 : f32
   // CHECK: 1.000000e+00 : f64
+  // CHECK: 1.000000e+00 : bf16
 
   int64_t indices[] = {0, 1};
   int64_t one = 1;
   MlirAttribute indicesAttr = mlirDenseElementsAttrInt64Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64), encoding),
       2, indices);
   MlirAttribute valuesAttr = mlirDenseElementsAttrFloatGet(
       mlirRankedTensorTypeGet(1, &one, mlirF32TypeGet(ctx), encoding), 1,