Add Tensor::tensor_id(): a per-object id safe to cache on, unlike &tensor

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<char>{proto_stream}, std::istreambuf_iterator<char>{}};
-  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<std::streamsize>(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<std::uintmax_t>(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<char>{proto_stream}, std::istreambuf_iterator<char>{});
+    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())) {

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<Symbol>
+  // 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);

onnx/common/tensor.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <cmath>
+#include <cstdint>
 #include <string>
 #include <utility>
 #include <vector>
@@ -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<std::pair<std::string, std::string>> 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<int64_t>& sizes() const {
     return sizes_;
   }

tests/cpp/tensor_test.cc

@@ -2,10 +2,13 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#include <algorithm>
 #include <cstdint>
 #include <limits>
 #include <stdexcept>
 #include <string>
+#include <utility>
+#include <vector>
 
 #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<int64_t>{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<int64_t>{5}));
+}
+
+TEST(TensorTest, TensorIdManyConstructionsAreAllDistinct) {
+  std::vector<uint64_t> ids;
+  ids.reserve(1000);
+  for (int i = 0; i < 1000; ++i) {
+    Tensor t;
+    ids.push_back(t.tensor_id());
+  }
+  std::vector<uint64_t> 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;