diff --git a/onnx/common/file_utils.h b/onnx/common/file_utils.h index 50d54541867..cca87985b51 100644 --- a/onnx/common/file_utils.h +++ b/onnx/common/file_utils.h @@ -26,8 +26,28 @@ void LoadProtoFromPath(const std::string& proto_path, T& proto) { if (!proto_stream.good()) { fail_check("Unable to open proto file: ", proto_path, ". Please check if it is a valid proto. "); } - std::string data{std::istreambuf_iterator{proto_stream}, std::istreambuf_iterator{}}; - if (!proto_stream.good()) { + // A single bulk read sized from the file's byte length, rather than + // istreambuf_iterator's byte-at-a-time copy (each increment pays a + // buffer-boundary check) -- on a large (hundreds of MB+) model file the + // difference is seconds, not microseconds. + std::error_code size_ec; + const std::uintmax_t file_size = std::filesystem::file_size(proto_u8_path, size_ec); + std::string data; + bool read_ok = false; + if (!size_ec) { + data.resize(file_size); + proto_stream.read(data.data(), static_cast(file_size)); + // read() may set eofbit alongside a fully successful read that consumes + // exactly to the end of the file, so check the actual byte count rather + // than the stream's good()/fail() flags. + read_ok = static_cast(proto_stream.gcount()) == file_size; + } else { + // Fall back to the iterator-based read if the size could not be + // determined (e.g. a non-regular file such as a pipe). + data.assign(std::istreambuf_iterator{proto_stream}, std::istreambuf_iterator{}); + read_ok = true; + } + if (!read_ok) { fail_check("Unable to read proto file: ", proto_path, ". Please check if it is a valid proto. "); } if (!ParseProtoFromBytes(&proto, data.c_str(), data.size())) { diff --git a/onnx/common/ir.h b/onnx/common/ir.h index 41c82844797..77eed3ad411 100644 --- a/onnx/common/ir.h +++ b/onnx/common/ir.h @@ -199,6 +199,23 @@ struct Attributes { names.push_back(a->name); return names; } + // Equivalent to checking whether attributeNames() contains a Symbol whose + // kindOf() is g or gs, but without attributeNames()'s std::vector + // allocation -- values_ already stores each attribute's kind directly, so + // this only needs to walk the (typically short) existing vector. Used by + // Graph::forEachNode's subgraph search (see ir.h's forSelfAndEachSubGraphImpl), + // which calls this once per node in the graph to find nodes with a nested + // subgraph; that search runs on every Value::uses()/setUniqueName() call, + // so avoiding an allocation here matters at scale. + bool hasSubgraphAttribute() const { + for (const auto& a : values_) { + const auto kind = a->kind(); + if (kind == AttributeKind::g || kind == AttributeKind::gs) { + return true; + } + } + return false; + } #define CREATE_ACCESSOR(Kind, method) \ Derived* method##_(Symbol name, Kind##Attr::ConstructorType v) { \ @@ -1261,6 +1278,15 @@ struct Graph final { fn(self); for (const auto& node_entry : self->all_nodes) { const Node* node = node_entry.first; + // hasSubgraphAttribute() is a cheap, non-allocating check; skip the + // allocating attributeNames() call entirely for the overwhelming + // majority of nodes (no g/gs attribute at all -- i.e. every node + // outside a Loop/If/Scan-style op). This traversal runs on every + // Value::uses()/setUniqueName() call, so avoiding that allocation on + // every one of a graph's nodes matters at scale. + if (!node->hasSubgraphAttribute()) { + continue; + } for (const auto& attr : node->attributeNames()) { if (node->kindOf(attr) == AttributeKind::g) { forSelfAndEachSubGraphImpl(node->g(attr).get(), fn); diff --git a/onnx/common/tensor.h b/onnx/common/tensor.h index 8a3892c38a7..e7d46d09ef5 100644 --- a/onnx/common/tensor.h +++ b/onnx/common/tensor.h @@ -8,6 +8,7 @@ #pragma once #include +#include #include #include #include @@ -20,6 +21,29 @@ namespace ONNX_NAMESPACE { struct Tensor final { private: + // A process-wide-unique id, freshly (re-)assigned by every constructor AND + // every assignment operator below -- never preserved across a copy, a + // move, or a reassignment. This makes it safe to key a cache off + // (tensor_id_) instead of (&tensor): a `Tensor*` can be freed and its + // memory reused by an unrelated, later Tensor (e.g. after + // Graph::eraseInitializer, or a Node attribute being replaced), which + // would silently alias an old cache entry onto the new tensor's different + // content if the cache were keyed by address. tensor_id_ can't collide + // this way: a fresh id is minted every time a Tensor's content is + // established or changed, so two live objects (or an old, freed one and a + // new one reusing its address) never share an id while their contents + // could differ. See onnxoptimizer/passes/tensor_content_hash.h's + // TensorContentDigest cache, the motivating consumer (onnxsim issue #633). + // + // Not atomic: Tensor construction is single-threaded throughout this + // codebase (matching TensorContentDigest's own cache, which is likewise + // an unsynchronized global). + static uint64_t NextTensorId() { + static uint64_t counter = 0; + return counter++; + } + uint64_t tensor_id_{NextTensorId()}; + bool is_segment_{false}; int64_t segment_begin_{0}; int64_t segment_end_{0}; @@ -41,7 +65,101 @@ struct Tensor final { std::vector> external_data_; ONNX_NAMESPACE::TensorProto_DataLocation data_location_{ONNX_NAMESPACE::TensorProto_DataLocation_DEFAULT}; + // Copies every field except tensor_id_, which each caller below sources + // independently (a fresh id via NextTensorId() for the constructors' + // member-initializer lists, this object's own existing id -- reassigned + // right after -- for the assignment operators). + void CopyFieldsFrom(const Tensor& other) { + is_segment_ = other.is_segment_; + segment_begin_ = other.segment_begin_; + segment_end_ = other.segment_end_; + has_name_ = other.has_name_; + name_ = other.name_; + elem_type_ = other.elem_type_; + sizes_ = other.sizes_; + float_data_ = other.float_data_; + double_data_ = other.double_data_; + int32_data_ = other.int32_data_; + int64_data_ = other.int64_data_; + uint64_data_ = other.uint64_data_; + string_data_ = other.string_data_; + is_raw_data_ = other.is_raw_data_; + raw_data_ = other.raw_data_; + external_data_ = other.external_data_; + data_location_ = other.data_location_; + } + // Every member below is individually moved from `other`, which is the + // correct and complete way to move-from an rvalue-reference parameter -- + // but cppcoreguidelines-rvalue-reference-param-not-moved only recognizes + // std::move(other) applied to the parameter itself, not std::move(other.x) + // on its members, so it flags this as a false positive. + // NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved) + void MoveFieldsFrom(Tensor&& other) { + is_segment_ = other.is_segment_; + segment_begin_ = other.segment_begin_; + segment_end_ = other.segment_end_; + has_name_ = other.has_name_; + name_ = std::move(other.name_); + elem_type_ = other.elem_type_; + sizes_ = std::move(other.sizes_); + float_data_ = std::move(other.float_data_); + double_data_ = std::move(other.double_data_); + int32_data_ = std::move(other.int32_data_); + int64_data_ = std::move(other.int64_data_); + uint64_data_ = std::move(other.uint64_data_); + string_data_ = std::move(other.string_data_); + is_raw_data_ = other.is_raw_data_; + raw_data_ = std::move(other.raw_data_); + external_data_ = std::move(other.external_data_); + data_location_ = other.data_location_; + } + public: + Tensor() = default; + // Explicit despite being a no-op: the copy/move constructor and + // copy/move assignment operator below are all user-declared (needed for + // tensor_id_'s semantics), which otherwise leaves the destructor as the + // sole implicit special member -- cppcoreguidelines-special-member-functions + // requires it be declared too for the rule-of-five to be unambiguous. + ~Tensor() = default; + // tensor_id_ deliberately omitted from these two constructors' behavior: + // it keeps its own default member initializer (a fresh NextTensorId()), + // never other's -- a copy/move-constructed Tensor is a logically distinct + // object from its source, even when byte-identical right now, since nothing + // stops either one from being separately reassigned or fed through + // graph.addInitializer() (which itself copies) afterward. + Tensor(const Tensor& other) { + CopyFieldsFrom(other); + } + Tensor(Tensor&& other) noexcept { + MoveFieldsFrom(std::move(other)); + } + Tensor& operator=(const Tensor& other) { + if (this != &other) { + CopyFieldsFrom(other); + // This object's content just changed, so any cache entry keyed on its + // previous tensor_id_ is now stale -- mint a fresh one rather than + // keeping this object's existing id (which would leave that stale + // entry silently reachable) or other's (which would collide with + // other's own, still-live id). + tensor_id_ = NextTensorId(); + } + return *this; + } + Tensor& operator=(Tensor&& other) noexcept { + if (this != &other) { + MoveFieldsFrom(std::move(other)); + tensor_id_ = NextTensorId(); + } + return *this; + } + + // See tensor_id_'s own comment above for the identity/uniqueness + // guarantee this provides. + uint64_t tensor_id() const { + return tensor_id_; + } + const std::vector& sizes() const { return sizes_; } diff --git a/tests/cpp/tensor_test.cc b/tests/cpp/tensor_test.cc index ce6cb4a8069..706b8a4f0a3 100644 --- a/tests/cpp/tensor_test.cc +++ b/tests/cpp/tensor_test.cc @@ -2,10 +2,13 @@ // // SPDX-License-Identifier: Apache-2.0 +#include #include #include #include #include +#include +#include #include "gtest/gtest.h" #include "onnx/common/assertions.h" @@ -51,6 +54,90 @@ TEST(TensorTest, SizeFromDimOverflowThrows) { #endif } +// tensor_id() backs onnxoptimizer's TensorContentDigest cache +// (onnxoptimizer/passes/tensor_content_hash.h), which relies on it never +// being shared between two Tensor objects whose content could independently +// diverge -- these tests cover the identity/uniqueness guarantee that +// invariant depends on. See tensor.h's tensor_id_ comment for the full +// rationale. +TEST(TensorTest, TensorIdDistinctAcrossDefaultConstruction) { + Tensor a; + Tensor b; + EXPECT_NE(a.tensor_id(), b.tensor_id()); +} + +TEST(TensorTest, TensorIdDistinctAcrossCopyConstruction) { + Tensor a; + a.sizes() = {1, 2, 3}; + Tensor b(a); + EXPECT_NE(a.tensor_id(), b.tensor_id()); + // Copying is a deep, independent copy: mutating one's content afterward + // (as several onnxoptimizer passes do -- fetch-by-copy, mutate the local + // copy, add it to the graph as new content) must not affect the other's. + b.sizes().push_back(4); + EXPECT_NE(a.sizes(), b.sizes()); +} + +TEST(TensorTest, TensorIdDistinctAcrossMoveConstruction) { + Tensor a; + a.sizes() = {1, 2, 3}; + const uint64_t a_id = a.tensor_id(); + Tensor b(std::move(a)); + EXPECT_NE(a_id, b.tensor_id()); + EXPECT_EQ(b.sizes(), (std::vector{1, 2, 3})); +} + +TEST(TensorTest, TensorIdRefreshedByCopyAssignment) { + Tensor a; + Tensor b; + const uint64_t b_id_before = b.tensor_id(); + b = a; + // b's content just changed (even if, as here, to something that happens to + // look the same) -- its id must move on from whatever it was before, so a + // cache entry keyed on the old id is never silently reattached to it. + EXPECT_NE(b_id_before, b.tensor_id()); + EXPECT_NE(a.tensor_id(), b.tensor_id()); +} + +TEST(TensorTest, TensorIdRefreshedByMoveAssignment) { + Tensor a; + Tensor b; + const uint64_t a_id = a.tensor_id(); + const uint64_t b_id_before = b.tensor_id(); + b = std::move(a); + EXPECT_NE(b_id_before, b.tensor_id()); + EXPECT_NE(a_id, b.tensor_id()); +} + +TEST(TensorTest, TensorIdSelfAssignmentIsANoop) { + Tensor a; + a.sizes() = {5}; + const uint64_t id_before = a.tensor_id(); +#if defined(__clang__) +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wself-assign-overloaded" +#endif + a = a; +#if defined(__clang__) +#pragma clang diagnostic pop +#endif + EXPECT_EQ(id_before, a.tensor_id()); + EXPECT_EQ(a.sizes(), (std::vector{5})); +} + +TEST(TensorTest, TensorIdManyConstructionsAreAllDistinct) { + std::vector ids; + ids.reserve(1000); + for (int i = 0; i < 1000; ++i) { + Tensor t; + ids.push_back(t.tensor_id()); + } + std::vector sorted_ids = ids; + std::sort(sorted_ids.begin(), sorted_ids.end()); + EXPECT_EQ(std::adjacent_find(sorted_ids.begin(), sorted_ids.end()), sorted_ids.end()) + << "expected all 1000 tensor_id()s to be distinct"; +} + TEST(TensorTest, ParseDataThrowsOnMisalignedRawData) { #ifndef ONNX_NO_EXCEPTIONS Tensor t;