[Mlir-commits] [mlir] [mlir][Target] Teach dense_resource conversion to LLVMIR Target (PR #78958)
Kunwar Grover
llvmlistbot at llvm.org
Mon Jan 22 11:03:19 PST 2024
https://github.com/Groverkss updated https://github.com/llvm/llvm-project/pull/78958
>From 977b7b25657cca7baa31e2f659c492fa03f7f046 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Mon, 22 Jan 2024 16:20:00 +0530
Subject: [PATCH 1/3] [mlir][Target] Teach LLVMIR target conversion for
dense_resource attr
---
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 91 ++++++++++++++++++++
1 file changed, 91 insertions(+)
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());
>From c4c78129c0ce094bc4a85aad0125fd1ecd8a47e8 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Mon, 22 Jan 2024 16:33:44 +0530
Subject: [PATCH 2/3] Add test
---
mlir/test/Target/LLVMIR/llvmir.mlir | 12 ++++++++++++
1 file changed, 12 insertions(+)
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
>From 7e2626fec37e3524c46ec5109db64820978d9ab4 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Tue, 23 Jan 2024 00:33:02 +0530
Subject: [PATCH 3/3] Address comments
---
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 7 +++++--
mlir/test/Target/LLVMIR/llvmir.mlir | 17 ++++++++++++++---
2 files changed, 19 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index a5f2400c0659669..3042f98c6ef2856 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -468,8 +468,11 @@ static llvm::Constant *convertDenseResourceElementsAttr(
if (type.getNumElements() == 0)
return nullptr;
- ArrayRef<char> rawData =
- denseResourceAttr.getRawHandle().getBlob()->getData();
+ AsmResourceBlob *blob = denseResourceAttr.getRawHandle().getBlob();
+ if (!blob)
+ return nullptr;
+
+ ArrayRef<char> rawData = blob->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
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index aedda4dfb7d0b3e..448aa3a5d85d79a 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -101,8 +101,18 @@ 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>
+// CHECK{LITERAL}: @dense_resource_tensor_constant = internal constant [5 x float] [float 0x3FCA034080000000, float 0xBFD0466300000000, float 0xBFD75DDF80000000, float 0xBFDE074F40000000, float 0x3FDDD3A1C0000000]
+llvm.mlir.global internal constant @dense_resource_tensor_constant(dense_resource<dense_resource_test_5xf32> : tensor<5xf32>) : !llvm.array<5 x f32>
+
+// CHECK{LITERAL}: @dense_resource_vector_constant = internal constant <5 x float> <float 0x3FCA034080000000, float 0xBFD0466300000000, float 0xBFD75DDF80000000, float 0xBFDE074F40000000, float 0x3FDDD3A1C0000000>
+llvm.mlir.global internal constant @dense_resource_vector_constant(dense_resource<dense_resource_test_5xf32> : vector<5xf32>) : vector<5xf32>
+
+
+// CHECK{LITERAL}: @dense_resource_multidim_tensor_constant = internal constant [1 x [2 x [2 x float]]] [[2 x [2 x float]] [[2 x float] [float 0x3FD6B46A80000000, float 0x3FD6781AC0000000], [2 x float] [float 0xBFB45A2AA0000000, float 0x3FD77A5CA0000000]]]
+llvm.mlir.global internal constant @dense_resource_multidim_tensor_constant(dense_resource<dense_resource_test_2x2xf32> : tensor<1x2x2xf32>) : !llvm.array<1 x !llvm.array<2 x !llvm.array<2 x f32>>>
+
+// CHECK{LITERAL}: @dense_resource_multidim_vector_constant = internal constant [1 x [2 x <2 x float>]] [[2 x <2 x float>] [<2 x float> <float 0x3FD6B46A80000000, float 0x3FD6781AC0000000>, <2 x float> <float 0xBFB45A2AA0000000, float 0x3FD77A5CA0000000>]]
+llvm.mlir.global internal constant @dense_resource_multidim_vector_constant(dense_resource<dense_resource_test_2x2xf32> : vector<1x2x2xf32>) : !llvm.array<1 x !llvm.array<2 x vector<2 x f32>>>
//
// Linkage attribute.
@@ -1584,7 +1594,8 @@ llvm.func @invokeLandingpad() -> i32 attributes { personality = @__gxx_personali
{-#
dialect_resources: {
builtin: {
- dense_resource_test: "0x08000000041A503E183382BEFCEEBABE7A3AF0BE0E9DEE3E"
+ dense_resource_test_5xf32: "0x08000000041A503E183382BEFCEEBABE7A3AF0BE0E9DEE3E",
+ dense_resource_test_2x2xf32: "0x0800000054A3B53ED6C0B33E55D1A2BDE5D2BB3E"
}
}
#-}
More information about the Mlir-commits
mailing list