use the same mmap for all tensors from the same external data file

src/Builder/FrontendDialectHelper.cpp

@@ -17,8 +17,6 @@
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Endian.h"
-#include "llvm/Support/MemoryBuffer.h"
-#include "llvm/Support/Path.h"
 #include "llvm/Support/SwapByteOrder.h"
 
 #include "src/Dialect/ONNX/ElementsAttr/Arrays.hpp"
@@ -33,47 +31,6 @@ namespace onnx_mlir {
 
 namespace {
 
-// 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.
-// The data is little endian encoded.
-// See https://github.com/onnx/onnx/blob/main/docs/ExternalData.md
-std::unique_ptr<llvm::MemoryBuffer> readExternalData_LE(
-    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");
-  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<int32_t> &data(
@@ -156,15 +113,15 @@ T swappedBytes(T x) {
 
 template <typename T>
 ElementsAttr createElementsAttrFromMemoryBuffer_LE(
-    RankedTensorType tensorType, std::unique_ptr<llvm::MemoryBuffer> membuf) {
+    RankedTensorType tensorType, const ExternalDataFileSlice &fileSlice) {
   MLIRContext *ctx = tensorType.getContext();
   assert(tensorType.getElementType() == toMlirType<T>(ctx));
   if constexpr (shouldSwapLEBytes<T>) {
-    ArrayRef<T> array = asArrayRef<T>(membuf->getBuffer());
+    ArrayRef<T> array = asArrayRef<T>(fileSlice.getBufferSlice());
     return createElmAttrFromArray<T>(tensorType, array, swappedBytes<T>);
   } else {
     return OnnxElementsAttrBuilder(ctx).fromMemoryBuffer(
-        tensorType, std::move(membuf));
+        tensorType, fileSlice.file, fileSlice.offset, fileSlice.length);
   }
 }
 
@@ -203,11 +160,12 @@ ElementsAttr createElmAttrFromProtoData(RankedTensorType tensorType,
 // Returns ElementsAttr with tp's data.
 template <typename T>
 ElementsAttr createElmAttr(RankedTensorType tensorType,
-    const onnx::TensorProto &tp, const std::string &externalDataDir) {
+    const onnx::TensorProto &tp,
+    const ExternalDataFileSlice *externalDataFileSlice) {
   if (tp.has_data_location() &&
       tp.data_location() == onnx::TensorProto::EXTERNAL) {
     return createElementsAttrFromMemoryBuffer_LE<T>(
-        tensorType, readExternalData_LE(externalDataDir, tp));
+        tensorType, *externalDataFileSlice);
   }
   if (tp.has_raw_data()) {
     return createElmAttrFromRawBytes_LE<T>(
@@ -236,7 +194,8 @@ ElementsAttr createStringElmAttr(
 } // namespace
 
 ElementsAttr onnxTensorProtoToElmAttr(MLIRContext *ctx,
-    const std::string &externalDataDir, const onnx::TensorProto &tp) {
+    const onnx::TensorProto &tp,
+    const ExternalDataFileSlice *externalDataFileSlice) {
   // Tensor dimensions.
   ArrayRef<int64_t> tensorDims(tp.dims().data(), tp.dims().size());
   if (tp.data_type() == onnx::TensorProto::STRING) {
@@ -249,7 +208,7 @@ ElementsAttr onnxTensorProtoToElmAttr(MLIRContext *ctx,
   auto tensorType = RankedTensorType::get(tensorDims, elmType);
   return dispatchByBType(btype, [&](auto btype) {
     using cpptype = CppType<btype>;
-    return createElmAttr<cpptype>(tensorType, tp, externalDataDir);
+    return createElmAttr<cpptype>(tensorType, tp, externalDataFileSlice);
   });
 }
 

src/Builder/FrontendDialectHelper.hpp

@@ -15,14 +15,26 @@
 #pragma once
 
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/MemoryBuffer.h"
 
 #include "onnx/onnx_pb.h"
 
-#include <string>
+#include <memory>
 
 namespace onnx_mlir {
 
+struct ExternalDataFileSlice {
+  std::shared_ptr<llvm::MemoryBuffer> file;
+  uint64_t offset;
+  uint64_t length;
+  llvm::StringRef getBufferSlice() const {
+    return file->getBuffer().substr(offset, length);
+  }
+};
+
 mlir::ElementsAttr onnxTensorProtoToElmAttr(mlir::MLIRContext *ctx,
-    const std::string &externalDataDir, const onnx::TensorProto &initializer);
+    const onnx::TensorProto &initializer,
+    const ExternalDataFileSlice *externalDataFileSlice = nullptr);
 
 } // namespace onnx_mlir

src/Builder/FrontendDialectTransformer.cpp

@@ -25,6 +25,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 
 #include "include/onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Builder/FrontendDialectTransformer.hpp"
@@ -142,6 +143,14 @@ void replaceAttrRefs(onnx::GraphProto &graph, const AttrMap &attr_map) {
 // End of copied code from third_party/onnx.
 // -------------------------------------------------------------------------- //
 
+// 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;
+}
+
 } // namespace
 
 namespace detail {
@@ -211,6 +220,66 @@ class FrontendGenImpl {
 
   ModelLocalFunctionsMap in_model_functions_;
 
+  using ExternalDataFiles =
+      std::unordered_map<std::string, std::shared_ptr<llvm::MemoryBuffer>>;
+
+  ExternalDataFiles externalDataFiles_;
+
+  const std::shared_ptr<llvm::MemoryBuffer> &mapExternalDataFile(
+      const std::string &location) {
+    auto [iter, inserted] = externalDataFiles_.try_emplace(location, nullptr);
+    if (inserted) {
+      StringRef dir = options_.externalDataDir;
+      SmallVector<char> pathVector(dir.begin(), dir.end());
+      llvm::sys::path::append(pathVector, location);
+      StringRef path(pathVector.data(), pathVector.size());
+      // Memory maps file (in most cases) or reads it into memory.
+      auto bufferOrError = llvm::MemoryBuffer::getFile(
+          path, /*IsText=*/false, /*RequiresNullTerminator=*/false);
+      if (std::error_code ec = bufferOrError.getError()) {
+        llvm::errs() << "Error " << ec.message() << " reading from file "
+                     << path << "\n";
+        llvm_unreachable("llvm::MemoryBuffer::getFile failed");
+      }
+      std::unique_ptr<llvm::MemoryBuffer> buffer =
+          std::move(bufferOrError.get());
+      assert(buffer->getBufferIdentifier() == path &&
+             "buffer identifier is file path");
+      iter->second = std::move(buffer);
+    }
+    return iter->second;
+  }
+
+  ExternalDataFileSlice readExternalData(const onnx::TensorProto &tp) {
+    assert(tp.has_data_location() &&
+           tp.data_location() == onnx::TensorProto::EXTERNAL &&
+           "tensor proto data must be external");
+    // MemoryBuffer uses -1 to mean infinity
+    constexpr uint64_t infiniteLength = -1;
+    std::string location;
+    uint64_t offset = 0;
+    uint64_t length = infiniteLength;
+    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");
+    const std::shared_ptr<llvm::MemoryBuffer> &file =
+        mapExternalDataFile(location);
+    uint64_t fileLength = file->getBufferSize();
+    assert(offset <= fileLength && "offset out of range");
+    if (length != infiniteLength)
+      assert(offset + length <= fileLength && "length out of range");
+    return {file, offset, length};
+  }
+
   Location UnknownLoc() const { return UnknownLoc::get(&context_); }
 
   Location ImportLoc(const onnx::NodeProto &node) {
@@ -264,8 +333,14 @@ class FrontendGenImpl {
   }
 
   Value ImportTensor(const onnx::TensorProto &tensor) {
-    mlir::ElementsAttr mlirAttr =
-        onnxTensorProtoToElmAttr(&context_, options_.externalDataDir, tensor);
+    mlir::ElementsAttr mlirAttr;
+    if (tensor.has_data_location() &&
+        tensor.data_location() == onnx::TensorProto::EXTERNAL) {
+      ExternalDataFileSlice fileSlice = readExternalData(tensor);
+      mlirAttr = onnxTensorProtoToElmAttr(&context_, tensor, &fileSlice);
+    } else {
+      mlirAttr = onnxTensorProtoToElmAttr(&context_, tensor);
+    }
     // Use the tensor name as Location.
     auto loc =
         NameLoc::get(builder_.getStringAttr("Initializer_" + tensor.name()));
@@ -385,10 +460,16 @@ class FrontendGenImpl {
       mlirAttr = builder_.getI64ArrayAttr(
           llvm::ArrayRef(attr.ints().data(), attr.ints().size()));
       break;
-    case onnx::AttributeProto::TENSOR:
-      mlirAttr = onnxTensorProtoToElmAttr(
-          &context_, options_.externalDataDir, attr.t());
-      break;
+    case onnx::AttributeProto::TENSOR: {
+      const onnx::TensorProto &tensor = attr.t();
+      if (tensor.has_data_location() &&
+          tensor.data_location() == onnx::TensorProto::EXTERNAL) {
+        ExternalDataFileSlice fileSlice = readExternalData(tensor);
+        mlirAttr = onnxTensorProtoToElmAttr(&context_, tensor, &fileSlice);
+      } else {
+        mlirAttr = onnxTensorProtoToElmAttr(&context_, tensor);
+      }
+    } break;
     case onnx::AttributeProto::STRINGS: {
       llvm::SmallVector<StringRef, 4> vectorStringRef;
       for (const auto &item : attr.strings()) {
@@ -1343,6 +1424,9 @@ class FrontendGenImpl {
     auto funcType = importGraph(graph, /*region=*/mainFunc.getBody(),
         /*op=*/mainFunc.getOperation(), /*useReturn=*/true);
     mainFunc.setType(funcType);
+    if (!externalDataFiles_.empty())
+      mainFunc->setAttr("external_data_dir",
+          builder_.getStringAttr(options_.externalDataDir));
 
     // Emit entry point op describing inference function signature.
     auto entryPoint = ONNXEntryPointOp::create(UnknownLoc(), mainFunc);