[mlir][Target] Teach dense_resource conversion to LLVMIR Target

[Groverkss]

This patch adds support for translating dense_resource attributes to LLVMIR Target.
The support added is similar to how DenseElementsAttr is handled, except we
don't need to handle splats.

Another possible way of doing this is adding iteration on dense_resource, but that is
non-trivial as DenseResourceAttr is not meant to be something you should directly
access. It has subclasses which you are supposed to use to iterate on it.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Kunwar Grover (Groverkss)

Changes

This patch adds support for translating dense_resource attributes to LLVMIR Target.

---

Full diff: https://github.com/llvm/llvm-project/pull/78958.diff

2 Files Affected:

• (modified) mlir/lib/Target/LLVMIR/ModuleTranslation.cpp (+91)
• (modified) mlir/test/Target/LLVMIR/llvmir.mlir (+12)

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 833b6b8e2f0fbf2..a5f2400c0659669 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -27,6 +27,7 @@
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
@@ -446,6 +447,90 @@ convertDenseElementsAttr(Location loc, DenseElementsAttr denseElementsAttr,
   return buildSequentialConstant(constantsRef, outerShape, llvmType, loc);
 }
 
+/// Convert a dense resource elements attribute to an LLVM IR constant using its
+/// raw data storage if possible. This supports elements attributes of tensor or
+/// vector type and avoids constructing separate objects for individual values
+/// of the innermost dimension. Constants for other dimensions are still
+/// constructed recursively. Returns null if constructing from raw data is not
+/// supported for this type, e.g., element type is not a power-of-two-sized
+/// primitive. Reports other errors at `loc`.
+static llvm::Constant *convertDenseResourceElementsAttr(
+    Location loc, DenseResourceElementsAttr denseResourceAttr,
+    llvm::Type *llvmType, const ModuleTranslation &moduleTranslation) {
+  if (!denseResourceAttr)
+    return nullptr;
+
+  llvm::Type *innermostLLVMType = getInnermostElementType(llvmType);
+  if (!llvm::ConstantDataSequential::isElementTypeCompatible(innermostLLVMType))
+    return nullptr;
+
+  ShapedType type = denseResourceAttr.getType();
+  if (type.getNumElements() == 0)
+    return nullptr;
+
+  ArrayRef<char> rawData =
+      denseResourceAttr.getRawHandle().getBlob()->getData();
+
+  // Check that the raw data size matches what is expected for the scalar size.
+  // TODO: in theory, we could repack the data here to keep constructing from
+  // raw data.
+  // TODO: we may also need to consider endianness when cross-compiling to an
+  // architecture where it is different.
+  int64_t numElements = denseResourceAttr.getType().getNumElements();
+  int64_t elementByteSize = rawData.size() / numElements;
+  if (8 * elementByteSize != innermostLLVMType->getScalarSizeInBits())
+    return nullptr;
+
+  // Compute the shape of all dimensions but the innermost. Note that the
+  // innermost dimension may be that of the vector element type.
+  bool hasVectorElementType = isa<VectorType>(type.getElementType());
+  int64_t numAggregates =
+      numElements / (hasVectorElementType
+                         ? 1
+                         : denseResourceAttr.getType().getShape().back());
+  ArrayRef<int64_t> outerShape = type.getShape();
+  if (!hasVectorElementType)
+    outerShape = outerShape.drop_back();
+
+  // Create a constructor for the innermost constant from a piece of raw data.
+  std::function<llvm::Constant *(StringRef)> buildCstData;
+  if (isa<TensorType>(type)) {
+    auto vectorElementType = dyn_cast<VectorType>(type.getElementType());
+    if (vectorElementType && vectorElementType.getRank() == 1) {
+      buildCstData = [&](StringRef data) {
+        return llvm::ConstantDataVector::getRaw(
+            data, vectorElementType.getShape().back(), innermostLLVMType);
+      };
+    } else if (!vectorElementType) {
+      buildCstData = [&](StringRef data) {
+        return llvm::ConstantDataArray::getRaw(data, type.getShape().back(),
+                                               innermostLLVMType);
+      };
+    }
+  } else if (isa<VectorType>(type)) {
+    buildCstData = [&](StringRef data) {
+      return llvm::ConstantDataVector::getRaw(data, type.getShape().back(),
+                                              innermostLLVMType);
+    };
+  }
+  if (!buildCstData)
+    return nullptr;
+
+  // Create innermost constants and defer to the default constant creation
+  // mechanism for other dimensions.
+  SmallVector<llvm::Constant *> constants;
+  int64_t aggregateSize = denseResourceAttr.getType().getShape().back() *
+                          (innermostLLVMType->getScalarSizeInBits() / 8);
+  constants.reserve(numAggregates);
+  for (unsigned i = 0; i < numAggregates; ++i) {
+    StringRef data(rawData.data() + i * aggregateSize, aggregateSize);
+    constants.push_back(buildCstData(data));
+  }
+
+  ArrayRef<llvm::Constant *> constantsRef = constants;
+  return buildSequentialConstant(constantsRef, outerShape, llvmType, loc);
+}
+
 /// Create an LLVM IR constant of `llvmType` from the MLIR attribute `attr`.
 /// This currently supports integer, floating point, splat and dense element
 /// attributes and combinations thereof. Also, an array attribute with two
@@ -546,6 +631,12 @@ llvm::Constant *mlir::LLVM::detail::getLLVMConstant(
     return result;
   }
 
+  if (llvm::Constant *result = convertDenseResourceElementsAttr(
+          loc, dyn_cast<DenseResourceElementsAttr>(attr), llvmType,
+          moduleTranslation)) {
+    return result;
+  }
+
   // Fall back to element-by-element construction otherwise.
   if (auto elementsAttr = dyn_cast<ElementsAttr>(attr)) {
     assert(elementsAttr.getShapedType().hasStaticShape());
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 961c9484446845a..aedda4dfb7d0b3e 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -101,6 +101,9 @@ llvm.mlir.global internal @dense_float_vector_3d(dense<[[[1.0, 2.0], [3.0, 4.0]]
 // CHECK{LITERAL}: @splat_float_vector_3d = internal global [2 x [2 x <2 x float>]] [[2 x <2 x float>] [<2 x float> <float 4.200000e+01, float 4.200000e+01>, <2 x float> <float 4.200000e+01, float 4.200000e+01>], [2 x <2 x float>] [<2 x float> <float 4.200000e+01, float 4.200000e+01>, <2 x float> <float 4.200000e+01, float 4.200000e+01>]]
 llvm.mlir.global internal @splat_float_vector_3d(dense<42.0> : vector<2x2x2xf32>) : !llvm.array<2 x !llvm.array<2 x vector<2xf32>>>
 
+// CHECK: @dense_resource_constant = internal constant [5 x float] [float 0x3FCA034080000000, float 0xBFD0466300000000, float 0xBFD75DDF80000000, float 0xBFDE074F40000000, float 0x3FDDD3A1C0000000]
+llvm.mlir.global internal constant @dense_resource_constant(dense_resource<dense_resource_test> : tensor<5xf32>) : !llvm.array<5 x f32>
+
 //
 // Linkage attribute.
 //
@@ -1577,6 +1580,15 @@ llvm.func @invokeLandingpad() -> i32 attributes { personality = @__gxx_personali
   llvm.invoke %9(%6, %0) to ^bb2 unwind ^bb1 vararg(!llvm.func<void (ptr, ...)>) : !llvm.ptr, (!llvm.ptr, i32) -> ()
 }
 
+// Resources are kept at end of file. New tests should be added above this.
+{-#
+  dialect_resources: {
+    builtin: {
+      dense_resource_test: "0x08000000041A503E183382BEFCEEBABE7A3AF0BE0E9DEE3E"
+    }
+  }
+#-}
+
 // -----
 
 llvm.func @foo() -> i8

[stellaraccident]

Thanks! Some small comments.

[stellaraccident]

Looks good to me except for Mehdi's comment about a failure case being checked.

[stellaraccident]

Much better - thank you for the extra care on errors and tests.

[stellaraccident]

@joker-eph your comments look resolved to me. Ok to land?

[joker-eph]

Yes, LG, thanks for checking!

[uweigand]

This seems to have broken the build bot on SystemZ: https://lab.llvm.org/buildbot/#/builders/199/builds/31136

This is likely another instance of the problem with parsing external resources on big-endian systems described here: #63469

[Groverkss]

This seems to have broken the build bot on SystemZ: https://lab.llvm.org/buildbot/#/builders/199/builds/31136

This is likely another instance of the problem with parsing external resources on big-endian systems described here: #63469

Thank you for letting me know about the issue.

I can disable these tests for big endianness for now, by sending a patch fixing it forward today. Do you think you could help me by checking if the new patch fixes the issue (Once I send the PR)? I don't have a big-endian machine, and I don't think the CI test it, so it would be hard for me to check it.

[uweigand]

I can disable these tests for big endianness for now, by sending a patch fixing it forward today. Do you think you could help me by checking if the new patch fixes the issue (Once I send the PR)? I don't have a big-endian machine, and I don't think the CI test it, so it would be hard for me to check it.

Sure, happy to do that.