Skip to content

Commit

Permalink
support external data (#1699)
Browse files Browse the repository at this point in the history 
* introduced ExternalDataReader

Signed-off-by: Soren Lassen <[email protected]>

* read external_data

Signed-off-by: Soren Lassen <[email protected]>

* simplified conversion from RepeatedField to SmallVector

Signed-off-by: Soren Lassen <[email protected]>

* skip copy to SmallVector unless needed

Signed-off-by: Soren Lassen <[email protected]>

* fix ArrayRef size

Signed-off-by: Soren Lassen <[email protected]>

* fix string type

Signed-off-by: Soren Lassen <[email protected]>

* remove ExternalDataReader

benchmarking showed that there was no benefit to caching opened memory buffers

Signed-off-by: Soren Lassen <[email protected]>

* support parsing uint onnx tensors

Signed-off-by: Soren Lassen <[email protected]>

* add --load_external_data option to onnx2json.py

Signed-off-by: Soren Lassen <[email protected]>

* onnxExternalizeData.py tool for constructing external data examples for testing

Signed-off-by: Soren Lassen <[email protected]>

* added zipmap.py script that generates zipmap.json parse lit test

Signed-off-by: Soren Lassen <[email protected]>

* added parse lit tests for raw/external/nonraw data

Signed-off-by: Soren Lassen <[email protected]>

* added comments in utils/testing to explain the models

Signed-off-by: Soren Lassen <[email protected]>

* simplify implementation

Signed-off-by: Soren Lassen <[email protected]>

* remove one more copy

Signed-off-by: Soren Lassen <[email protected]>

* remove yet another copy

Signed-off-by: Soren Lassen <[email protected]>

* fix little-endian logic

Signed-off-by: Soren Lassen <[email protected]>

* small refactor to prepare for future float16 support

Signed-off-by: Soren Lassen <[email protected]>

Signed-off-by: Soren Lassen <[email protected]>
Co-authored-by: Philip Lassen <[email protected]>
  • Loading branch information
@sorenlassen @philass
sorenlassen and philass authored Sep 20, 2022
1 parent 8f0fff3 commit 6d68373
Show file tree
Hide file tree
Showing 14 changed files with 1,812 additions and 218 deletions.
309 changes: 140 additions & 169 deletions src/Builder/FrontendDialectHelper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,232 +11,203 @@
// Helper methods for handling input ONNX models.
//
//===----------------------------------------------------------------------===//

#include "llvm/ADT/SmallVector.h"
#include <llvm/Support/Endian.h>
#include <llvm/Support/SwapByteOrder.h>
#include "llvm/Support/Endian.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/SwapByteOrder.h"

#include "src/Builder/FrontendDialectHelper.hpp"

namespace onnx_mlir {

// Parses unsigned number.
size_t parseOffsetOrLength(const std::string &value) {
char *end = nullptr;
size_t offsetOrLength = strtoull(value.c_str(), &end, 0);
assert(end != value.c_str() && "failed to parse offset or length");
return offsetOrLength;
}

// Reads external data from file location specified in tensor proto.
// See https://github.com/onnx/onnx/blob/main/docs/ExternalData.md
std::unique_ptr<llvm::MemoryBuffer> readExternalData(
const std::string &externalDataDir, const onnx::TensorProto &tp) {
std::string location;
uint64_t offset = 0;
uint64_t length = -1; // MemoryBuffer uses -1 to mean infinity
for (const onnx::StringStringEntryProto &entry : tp.external_data()) {
assert(entry.has_key() && "external_data entry must have key");
assert(entry.has_value() && "external_data entry must have value");
if (entry.key() == "location") {
location = entry.value();
} else if (entry.key() == "offset") {
offset = parseOffsetOrLength(entry.value());
} else if (entry.key() == "length") {
length = parseOffsetOrLength(entry.value());
}
}
assert(!location.empty() && "missing external data location");
llvm::SmallVector<char> path(externalDataDir.begin(), externalDataDir.end());
llvm::sys::path::append(path, location);
auto bufferOrError = llvm::MemoryBuffer::getFileSlice(
path, length, offset, /*IsVolatile=*/false);
if (std::error_code ec = bufferOrError.getError()) {
std::string pathStr(path.data(), path.size());
llvm::errs() << "Error " << ec.message() << " reading from file " << pathStr
<< ", offset=" << offset << ", length=" << length << "\n";
llvm_unreachable("llvm::MemoryBuffer::getFileSlice failed");
}
return std::move(bufferOrError.get());
}

template <typename T>
struct TransformValueToONNXData {
static const google::protobuf::RepeatedField<T> data(
onnx::TensorProto initializer) {
return google::protobuf::RepeatedField<T>();
static const google::protobuf::RepeatedField<int32_t> &data(
const onnx::TensorProto &tp) {
// int32_data is used for int32, uint8, int8, uint16, int16, bool
return tp.int32_data();
}
};

template <>
struct TransformValueToONNXData<double> {
static const google::protobuf::RepeatedField<double> data(
onnx::TensorProto initializer) {
return initializer.double_data();
static const google::protobuf::RepeatedField<double> &data(
const onnx::TensorProto &tp) {
return tp.double_data();
}
};

template <>
struct TransformValueToONNXData<float> {
static const google::protobuf::RepeatedField<float> data(
onnx::TensorProto initializer) {
return initializer.float_data();
}
};

template <>
struct TransformValueToONNXData<int16_t> {
static const google::protobuf::RepeatedField<int32_t> data(
onnx::TensorProto initializer) {
return initializer.int32_data();
}
};

template <>
struct TransformValueToONNXData<int32_t> {
static const google::protobuf::RepeatedField<int32_t> data(
onnx::TensorProto initializer) {
return initializer.int32_data();
static const google::protobuf::RepeatedField<float> &data(
const onnx::TensorProto &tp) {
return tp.float_data();
}
};

template <>
struct TransformValueToONNXData<int64_t> {
static const google::protobuf::RepeatedField<int64_t> data(
onnx::TensorProto initializer) {
return initializer.int64_data();
static const google::protobuf::RepeatedField<int64_t> &data(
const onnx::TensorProto &tp) {
return tp.int64_data();
}
};

template <>
struct TransformValueToONNXData<uint8_t> {
static const google::protobuf::RepeatedField<int32_t> data(
onnx::TensorProto initializer) {
return initializer.int32_data();
struct TransformValueToONNXData<uint32_t> {
static const google::protobuf::RepeatedField<uint64_t> &data(
const onnx::TensorProto &tp) {
return tp.uint64_data();
}
};

template <>
struct TransformValueToONNXData<int8_t> {
static const google::protobuf::RepeatedField<int32_t> data(
onnx::TensorProto initializer) {
return initializer.int32_data();
struct TransformValueToONNXData<uint64_t> {
static const google::protobuf::RepeatedField<uint64_t> &data(
const onnx::TensorProto &tp) {
return tp.uint64_data();
}
};

template <>
struct TransformValueToONNXData<bool> {
static const google::protobuf::RepeatedField<int32_t> data(
onnx::TensorProto initializer) {
return initializer.int32_data();
}
};
template <typename T>
using DenseElementsAttrBuilder =
llvm::function_ref<mlir::DenseElementsAttr(llvm::ArrayRef<T>)>;

// Helper method for constructing an array attribute from a model input.
// Returns DenseElementsAttr with tp's data.
template <typename T>
std::vector<T> CreateArrayAttribute(onnx::TensorProto initializer) {
// Return an emtpy array if the raw data is empty.
if (initializer.dims().size() == 1 && initializer.dims().data()[0] == 0)
return std::vector<T>();

size_t size;
if (initializer.raw_data().size()) {
// Copy & take care of endianness.
std::vector<char> byteInitializer;
std::copy(initializer.raw_data().begin(), initializer.raw_data().end(),
back_inserter(byteInitializer));
size = initializer.raw_data().size() / sizeof(T);
T *arrayPtr = reinterpret_cast<T *>(&byteInitializer[0]);
auto array = std::vector<T>(arrayPtr, arrayPtr + size);
mlir::DenseElementsAttr createDenseElmAttr(const std::string &externalDataDir,
const onnx::TensorProto &tp, DenseElementsAttrBuilder<T> denseBuilder) {
std::unique_ptr<llvm::MemoryBuffer> externalData =
(tp.has_data_location() &&
tp.data_location() == onnx::TensorProto::EXTERNAL)
? readExternalData(externalDataDir, tp)
: nullptr;
if (externalData || tp.has_raw_data()) {
llvm::StringRef buffer = externalData ? externalData->getBuffer()
: llvm::StringRef(tp.raw_data());
size_t size = buffer.size() / sizeof(T);
llvm::ArrayRef<T> arrayRef(
reinterpret_cast<T const *>(buffer.data()), size);
// Perform byte swap if system endianness is BE.
// ONNX tensor content raw data is always in LE.
if (llvm::support::endian::system_endianness() !=
llvm::support::endianness::little)
for (size_t i = 0; i < array.size(); i++)
llvm::sys::swapByteOrder<T>(array[i]);

return array;
if (sizeof(T) > 1 && llvm::support::endian::system_endianness() !=
llvm::support::endianness::little) {
llvm::SmallVector<T> vector;
vector.reserve(size);
for (T x : arrayRef) {
vector.push_back(llvm::sys::getSwappedBytes(x));
}
return denseBuilder(llvm::makeArrayRef(vector));
} else {
// No need to take care of endianness.
return denseBuilder(arrayRef);
}
} else {
// Not raw, no need to take care of endianness.
const auto &data = TransformValueToONNXData<T>::data(tp);
// Access data directly via ArrayRef if same size as T,
// or copy into correctly typed SmallVector otherwise
// because DenseElementsAttr needs argument type of the correct bitwidth.
typedef typename std::conditional<sizeof(T) == sizeof(data[0]),
llvm::ArrayRef<T>, llvm::SmallVector<T>>::type ArrayRefOrSmallVector;
ArrayRefOrSmallVector array(data.begin(), data.end());
return denseBuilder(llvm::makeArrayRef(array));
}

// Copy, no need to take care of endianness.
auto data = TransformValueToONNXData<T>::data(initializer);
size = data.size();
return std::vector<T>(&data[0], &data[0] + size);
}

template <>
std::vector<bool> CreateArrayAttribute<bool>(onnx::TensorProto initializer) {
// Copy, no need to take care of endianness.
if (initializer.raw_data().size()) {
std::vector<bool> bitInitializer;
std::copy(initializer.raw_data().begin(), initializer.raw_data().end(),
back_inserter(bitInitializer));
return bitInitializer;
}

auto data = TransformValueToONNXData<bool>::data(initializer);
return std::vector<bool>(&data[0], &data[0] + data.size());
}

mlir::Value EmitInitializerForInputTensor(mlir::Location loc,
mlir::OpBuilder &builder, const onnx::TensorProto &initializer) {
mlir::OpBuilder &builder, const std::string &externalDataDir,
const onnx::TensorProto &initializer) {
// Return none if the initializer is an empty tensor, e.g tensor<0xf32>.
llvm::ArrayRef<int64_t> tensorDims(
initializer.dims().data(), initializer.dims().size());
if (tensorDims.size() == 1 && tensorDims[0] == 0)
return builder.create<mlir::ONNXNoneOp>(
loc, builder.getNoneType(), builder.getUnitAttr());

// Emit ConstantOp and record the mapping between the input and
// the constant value.
// Create value attribute.
mlir::DenseElementsAttr denseElmAttr =
onnxTensorProtoToDenseElmAttr(builder, initializer);

// Create ConstantOp for dense array.
onnxTensorProtoToDenseElmAttr(builder, externalDataDir, initializer);
return builder.create<mlir::ONNXConstantOp>(loc, nullptr, denseElmAttr);
}

mlir::DenseElementsAttr onnxTensorProtoToDenseElmAttr(
mlir::OpBuilder &builder, const onnx::TensorProto &initializer) {
mlir::DenseElementsAttr onnxTensorProtoToDenseElmAttr(mlir::OpBuilder &builder,
const std::string &externalDataDir, const onnx::TensorProto &tp) {
// Tensor dimensions.
llvm::ArrayRef<int64_t> tensorDims(
initializer.dims().data(), initializer.dims().size());
mlir::DenseElementsAttr denseElmAttr;
switch (initializer.data_type()) {
case (onnx::TensorProto::FLOAT): {
const auto &arrayAttrInitializer = CreateArrayAttribute<float>(initializer);
auto elmType = builder.getF32Type();
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::DOUBLE): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<double>(initializer);
auto elmType = builder.getF64Type();
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::INT8): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<int8_t>(initializer);
auto elmType = builder.getIntegerType(8);
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::UINT8): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<uint8_t>(initializer);
auto elmType = builder.getIntegerType(8, false);
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::INT16): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<int16_t>(initializer);
auto elmType = builder.getIntegerType(16);
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::INT32): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<int32_t>(initializer);
auto elmType = builder.getIntegerType(32);
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::INT64): {
const auto &arrayAttrInitializer =
CreateArrayAttribute<int64_t>(initializer);
auto elmType = builder.getIntegerType(64);
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::makeArrayRef(arrayAttrInitializer));
break;
}
case (onnx::TensorProto::BOOL): {
const auto &data = CreateArrayAttribute<bool>(initializer);
auto elmType = builder.getI1Type();
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
denseElmAttr = mlir::DenseElementsAttr::get(
tensorType, llvm::SmallVector<bool, 64>(data.begin(), data.end()));
break;
}
llvm::ArrayRef<int64_t> tensorDims(tp.dims().data(), tp.dims().size());
mlir::Type elmType = convertONNXTypeToMLIRType(
builder, (onnx::TensorProto_DataType)tp.data_type());
auto tensorType = mlir::RankedTensorType::get(tensorDims, elmType);
auto denseBuilder = [tensorType](auto arrayRef) {
return mlir::DenseElementsAttr::get(tensorType, arrayRef);
};
switch (tp.data_type()) {
case (onnx::TensorProto::FLOAT):
return createDenseElmAttr<float>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::DOUBLE):
return createDenseElmAttr<double>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::INT8):
return createDenseElmAttr<int8_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::UINT8):
return createDenseElmAttr<uint8_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::INT16):
return createDenseElmAttr<int16_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::UINT16):
return createDenseElmAttr<uint16_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::INT32):
return createDenseElmAttr<int32_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::UINT32):
return createDenseElmAttr<uint32_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::INT64):
return createDenseElmAttr<int64_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::UINT64):
return createDenseElmAttr<uint64_t>(externalDataDir, tp, denseBuilder);
case (onnx::TensorProto::BOOL):
return createDenseElmAttr<bool>(externalDataDir, tp, denseBuilder);
default:
llvm_unreachable(
"Failed to import ONNX TensorProto due to unsupported data types.");
}
return denseElmAttr;
}
} // namespace onnx_mlir
7 changes: 4 additions & 3 deletions src/Builder/FrontendDialectHelper.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,10 @@
namespace onnx_mlir {

mlir::Value EmitInitializerForInputTensor(mlir::Location loc,
mlir::OpBuilder &builder, const onnx::TensorProto &initializer);
mlir::OpBuilder &builder, const std::string &externalDataDir,
const onnx::TensorProto &initializer);

mlir::DenseElementsAttr onnxTensorProtoToDenseElmAttr(
mlir::OpBuilder &builder, const onnx::TensorProto &initializer);
mlir::DenseElementsAttr onnxTensorProtoToDenseElmAttr(mlir::OpBuilder &builder,
const std::string &externalDataDir, const onnx::TensorProto &initializer);

} // namespace onnx_mlir
Loading

0 comments on commit 6d68373

Please sign in to comment.