[ET][native] Read bundled samples in C++

backends/native/runtime/CMakeLists.txt

@@ -45,8 +45,10 @@ add_library(native_reference_gguf_kquant utils/ReferenceGgufKQuant.cpp)
 target_link_libraries(native_reference_gguf_kquant PRIVATE Threads::Threads)
 
 add_library(
-  native_package deserialize/OwnedBytes.cpp deserialize/Package.cpp
-                 deserialize/SafeTensorsReader.cpp deserialize/ZipReader.cpp
+  native_package
+  deserialize/OwnedBytes.cpp deserialize/Package.cpp
+  deserialize/SafeTensorsReader.cpp deserialize/Samples.cpp
+  deserialize/ZipReader.cpp
 )
 target_link_libraries(
   native_package

backends/native/runtime/deserialize/Samples.cpp

@@ -0,0 +1,274 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <executorch/backends/native/runtime/deserialize/Samples.h>
+
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <stdexcept>
+#include <string_view>
+#include <unordered_set>
+#include <utility>
+
+#include <executorch/backends/native/runtime/deserialize/ByteSpan.h>
+#include <executorch/backends/native/runtime/deserialize/DeserializeError.h>
+#include <executorch/backends/native/runtime/deserialize/Json.h>
+#include <executorch/backends/native/runtime/deserialize/Limits.h>
+#include <executorch/backends/native/runtime/deserialize/SafeTensorsReader.h>
+
+namespace ptn {
+namespace {
+
+bool is_ssa_name(std::string_view name) {
+  return !name.empty() && !std::isdigit(static_cast<unsigned char>(name[0])) &&
+      std::all_of(name.begin(), name.end(), [](char c) {
+        return std::isalnum(static_cast<unsigned char>(c)) || c == '_';
+      });
+}
+
+bool is_set_name(std::string_view name) {
+  return !name.empty() && std::all_of(name.begin(), name.end(), [](char c) {
+    return std::isalnum(static_cast<unsigned char>(c)) || c == '_' || c == '-';
+  });
+}
+
+bool is_digits(std::string_view text) {
+  return !text.empty() && (text == "0" || text[0] != '0') &&
+      std::all_of(text.begin(), text.end(), [](char c) {
+        return std::isdigit(static_cast<unsigned char>(c));
+      });
+}
+
+// "MAJOR.MINOR" with major 1; a newer minor only adds fields.
+bool is_supported_version(const Json& version) {
+  if (!version.is_string()) {
+    return false;
+  }
+  const std::string_view text = version.get_ref<const std::string&>();
+  const size_t dot = text.find('.');
+  return dot != std::string_view::npos && text.substr(0, dot) == "1" &&
+      is_digits(text.substr(dot + 1));
+}
+
+// A newer minor may add fields that reference tensors this reader does not
+// know.
+bool is_newer_minor(const Json& version) {
+  const std::string_view text = version.get_ref<const std::string&>();
+  return text.substr(text.find('.') + 1) != "0";
+}
+
+// Rethrows the active exception prefixed with `where`, keeping the type of
+// limit errors.
+[[noreturn]] void rethrow_in(const std::string& where) {
+  try {
+    throw;
+  } catch (const ResourceLimitError& error) {
+    throw ResourceLimitError(where + ": " + error.what());
+  } catch (const std::exception& error) {
+    throw std::runtime_error(where + ": " + error.what());
+  }
+}
+
+class SetDecoder {
+ private:
+  std::string where_;
+  std::vector<uint8_t> blob_;
+  SafeTensorsReader tensors_;
+  size_t data_start_ = 0;
+  std::unordered_set<std::string> referenced_;
+
+ public:
+  SetDecoder(std::string where, std::vector<uint8_t> blob)
+      : where_(std::move(where)),
+        blob_(std::move(blob)),
+        tensors_(SafeTensorsReader::open(blob_)) {
+    data_start_ = SafeTensorsReader::kLengthPrefixSize +
+        SafeTensorsReader::header_size(ByteSpan(blob_).first(
+            SafeTensorsReader::kLengthPrefixSize));
+  }
+
+  [[noreturn]] void fail(const std::string& message) const {
+    throw std::runtime_error(where_ + ": " + message);
+  }
+
+  std::optional<Tolerance> tol(
+      const Json& object,
+      std::optional<Tolerance> fallback) const {
+    const auto it = object.find("tol");
+    if (it == object.end()) {
+      return fallback;
+    }
+    if (!it->is_array() || it->size() != 2) {
+      fail("tol must be [atol, rtol]");
+    }
+    for (const Json& t : *it) {
+      if (!t.is_number() || !std::isfinite(t.get<double>()) ||
+          t.get<double>() < 0) {
+        fail("tol must be two finite numbers >= 0");
+      }
+    }
+    return Tolerance{(*it)[0].get<double>(), (*it)[1].get<double>()};
+  }
+
+  SampleTensor tensor(const Json& value, const std::string& key) {
+    if (!value.is_string()) {
+      fail(key + " must be a tensor; scalar and list values are not supported");
+    }
+    const std::string& name = value.get_ref<const std::string&>();
+    if (name != key) {
+      fail("key " + name + " must be " + key);
+    }
+    const TensorEntry* entry = tensors_.find(name);
+    if (entry == nullptr) {
+      fail("no tensor named " + name);
+    }
+    referenced_.insert(name);
+    const uint8_t* begin = blob_.data() + data_start_ + entry->offset;
+    return {entry->dtype, entry->sizes, {begin, begin + entry->nbytes}};
+  }
+
+  std::vector<SampleTensor> tensors(
+      const Json& values,
+      const std::string& key) {
+    if (!values.is_array()) {
+      fail(key + " must be a list");
+    }
+    std::vector<SampleTensor> result;
+    result.reserve(values.size());
+    for (size_t i = 0; i < values.size(); ++i) {
+      result.push_back(tensor(values[i], key + "[" + std::to_string(i) + "]"));
+    }
+    return result;
+  }
+
+  SampleCall call(
+      const Json& entry,
+      const std::string& at,
+      std::optional<Tolerance> fallback) {
+    if (!entry.is_object() || !entry.contains("in")) {
+      fail(at + " must be an object with an 'in' list");
+    }
+    const auto method = entry.find("method");
+    if (method == entry.end() || !method->is_string() ||
+        method->get_ref<const std::string&>().empty()) {
+      fail(at + ".method must be a non-empty string");
+    }
+    SampleCall result;
+    result.method = method->get<std::string>();
+    result.inputs = tensors(entry["in"], at + ".in");
+    if (entry.contains("out")) {
+      result.outputs = tensors(entry["out"], at + ".out");
+    }
+    if (entry.contains("intermediates")) {
+      const Json& intermediates = entry["intermediates"];
+      if (!intermediates.is_object()) {
+        fail(at + ".intermediates must be an object");
+      }
+      const std::string key_prefix = at + ".intermediates.";
+      for (const auto& [name, key] : intermediates.items()) {
+        if (!is_ssa_name(name)) {
+          fail(std::string("bad intermediate ")
+                   .append(name)
+                   .append(" in ")
+                   .append(at));
+        }
+        result.intermediates.emplace(name, tensor(key, key_prefix + name));
+      }
+    }
+    result.tol = tol(entry, fallback);
+    return result;
+  }
+
+  void check_all_referenced() const {
+    for (const std::string& name : tensors_.names()) {
+      if (!referenced_.contains(name)) {
+        fail("unreferenced tensor " + name);
+      }
+    }
+  }
+};
+
+} // namespace
+
+// cppcheck-suppress unusedFunction
+SampleSet read_sample_set(const ZipReader& zip, const std::string& set) {
+  if (!is_set_name(set)) {
+    throw std::runtime_error(
+        "samples: set name '" + set + "' must match [A-Za-z0-9_-]+");
+  }
+  const std::string prefix = "samples/" + set + "/";
+  const std::string where = "samples: " + prefix;
+  const std::optional<size_t> index_size =
+      zip.member_size(prefix + "index.json");
+  if (!index_size || !zip.member_size(prefix + "data.safetensors")) {
+    throw std::runtime_error(
+        where + ": missing index.json or data.safetensors");
+  }
+  if (*index_size > detail::kMaxJsonBytes) {
+    throw ResourceLimitError(where + ": index.json exceeds size limit");
+  }
+  const std::vector<uint8_t> index_bytes = zip.read(prefix + "index.json");
+  Json index;
+  try {
+    index = parse_json(std::string_view(
+        reinterpret_cast<const char*>(index_bytes.data()), index_bytes.size()));
+  } catch (...) {
+    rethrow_in(where + ": invalid index.json");
+  }
+  if (!index.is_object()) {
+    throw std::runtime_error(where + ": index.json must contain a JSON object");
+  }
+  const auto version = index.find("version");
+  if (version == index.end() || !is_supported_version(*version)) {
+    throw UnsupportedVersionError(
+        where + ": unsupported version, expected 1.0");
+  }
+  std::optional<SetDecoder> decoder;
+  try {
+    decoder.emplace(where, zip.read(prefix + "data.safetensors"));
+  } catch (...) {
+    rethrow_in(where);
+  }
+
+  SampleSet result;
+  for (const auto& [field, value] :
+       {std::pair{"description", &result.description},
+        std::pair{"source", &result.source}}) {
+    if (index.contains(field)) {
+      if (!index[field].is_string()) {
+        decoder->fail(std::string(field) + " must be a string");
+      }
+      *value = index[field].get<std::string>();
+    }
+  }
+  const auto entries = index.find("samples");
+  if (entries == index.end() || !entries->is_array() || entries->empty()) {
+    decoder->fail("samples must be a non-empty list");
+  }
+  const std::optional<Tolerance> set_tol = decoder->tol(index, std::nullopt);
+  for (size_t i = 0; i < entries->size(); ++i) {
+    const Json& entry = (*entries)[i];
+    const std::string at = "samples[" + std::to_string(i) + "]";
+    if (!entry.is_object() || !entry.contains("calls") ||
+        !entry["calls"].is_array() || entry["calls"].empty()) {
+      decoder->fail(at + " must be an object with a non-empty 'calls' list");
+    }
+    const std::optional<Tolerance> sample_tol = decoder->tol(entry, set_tol);
+    Sample& sample = result.samples.emplace_back();
+    const Json& calls = entry["calls"];
+    for (size_t j = 0; j < calls.size(); ++j) {
+      sample.calls.push_back(decoder->call(
+          calls[j], at + ".calls[" + std::to_string(j) + "]", sample_tol));
+    }
+  }
+  if (!is_newer_minor(*version)) {
+    decoder->check_all_referenced();
+  }
+  return result;
+}
+
+} // namespace ptn

backends/native/runtime/deserialize/Samples.h

@@ -0,0 +1,61 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <executorch/backends/native/runtime/deserialize/ZipReader.h>
+#include <executorch/backends/native/runtime/graph/ScalarType.h>
+
+namespace ptn {
+
+// A recorded tensor, copied out of the package as contiguous little-endian
+// bytes.
+struct SampleTensor {
+  ScalarType dtype = kFloat;
+  std::vector<int64_t> sizes;
+  std::vector<uint8_t> bytes;
+};
+
+struct Tolerance {
+  double atol = 0.0;
+  double rtol = 0.0;
+};
+
+struct SampleCall {
+  std::string method;
+  std::vector<SampleTensor> inputs;
+  // Absent when the producer recorded no outputs.
+  std::optional<std::vector<SampleTensor>> outputs;
+  std::map<std::string, SampleTensor> intermediates;
+  // The call's tolerance, else its sample's, else its set's.
+  std::optional<Tolerance> tol;
+};
+
+// Calls run in order from the program's initial state, with no reset between.
+struct Sample {
+  std::vector<SampleCall> calls;
+};
+
+struct SampleSet {
+  std::string description;
+  std::string source;
+  std::vector<Sample> samples;
+};
+
+// Reads samples/<set>/, specified in backends/native/README.md. Every in and
+// out value must be a real-valued tensor; scalar, list and complex values are
+// not supported. Throws std::runtime_error if the set is absent or violates the
+// format: UnsupportedVersionError for another major version, and
+// ResourceLimitError when it exceeds a reader limit.
+SampleSet read_sample_set(const ZipReader& zip, const std::string& set);
+
+} // namespace ptn

backends/native/runtime/deserialize/targets.bzl

@@ -105,6 +105,24 @@ def define_common_targets():
         ],
         visibility = ["PUBLIC"],
     )
+    # Recorded method calls bundled under samples/ in a .ptn package.
+    runtime.cxx_library(
+        name = "samples",
+        srcs = ["Samples.cpp"],
+        exported_headers = ["Samples.h"],
+        exported_deps = [
+            ":zip_reader",
+            "//executorch/backends/native/runtime/graph:scalar_type",
+        ],
+        deps = [
+            ":byte_span",
+            ":deserialize_error",
+            ":json",
+            ":limits",
+            ":safetensors_reader",
+        ],
+        visibility = ["PUBLIC"],
+    )
     runtime.cxx_test(
         name = "reader_bounds_test",
         srcs = ["test/ReaderBoundsTest.cpp"],

backends/native/test/CMakeLists.txt

@@ -11,6 +11,7 @@ add_executable(
   ../runtime/deserialize/test/ValidationTest.cpp
   ../runtime/test/MethodMetaTest.cpp
   ../runtime/test/ProgramVersionTest.cpp
+  runtime/deserialize/test_samples.cpp
   runtime/graph/test_graph.cpp
   runtime/graph/test_graph_utils.cpp
   runtime/graph/test_memory_planning.cpp

backends/native/test/runtime/deserialize/targets.bzl

@@ -28,6 +28,16 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_test(
+        name = "samples_test",
+        srcs = ["test_samples.cpp"],
+        deps = [
+            "//executorch/backends/native/runtime/deserialize:deserialize_error",
+            "//executorch/backends/native/runtime/deserialize:package_test_data",
+            "//executorch/backends/native/runtime/deserialize:samples",
+        ],
+    )
+
     runtime.cxx_test(
         name = "zip_reader_test",
         srcs = ["test_zip_reader.cpp"],