[ET][native] Bundle sample method calls in .ptn packages

backends/native/README.md

@@ -41,6 +41,94 @@ A `.ptn` is an uncompressed zip with fixed entry names:
 | `program.safetensors` | Constants. Byte-identical immutable tensors are stored once. |
 | `aliases.json` | Deduplicated key to owner key; present only when needed. |
 
+## Samples (optional)
+
+A `.ptn` may carry recorded method calls: inputs, and optionally the outputs and
+intermediate values they produced. A sample is an ordered sequence of calls, each
+naming the method it runs, so a stateful sequence such as an LLM's prefill
+followed by decode steps can be checked whether it uses one method or several.
+Samples serve reference checks, layer-by-layer numerical debugging, calibration
+and benchmarking. Program loaders ignore them.
+
+```python
+from executorch.exir.native import Call, read_samples, Sample, SampleSet, to_native, write_samples
+
+to_native(ep).save("model.ptn")
+golden = SampleSet([Sample([Call("forward", [x], [model(x)])])])
+write_samples("model.ptn", "golden", golden)
+read_samples("model.ptn")       # {"golden": SampleSet(...)}
+```
+
+`write_samples` appends a set to any existing `.ptn`, not only one this process
+saved, and replaces the file only once the set is complete. `read_samples`
+validates every set against the rules below, except that each call's method
+exists in the program and its intermediate names are values of that method.
+
+### Layout
+
+```
+samples/<set>/index.json
+samples/<set>/data.safetensors
+```
+
+- `<set>` is an identifier with no meaning to the format, matching
+  `[A-Za-z0-9_-]+`.
+- Each set has exactly these two members, and nothing else lives under `samples/`.
+- Members are stored (uncompressed). The `data.safetensors` data section starts
+  64-byte aligned in the file, like `program.safetensors`.
+- Adding a set only appends members; removing one means rewriting the zip.
+
+### `index.json`
+
+```
+{ "version":     "1.0",                required, "MAJOR.MINOR"
+  "description": str,                  optional
+  "source":      str,                  optional, producer of the outputs
+  "tol":         [atol, rtol],         optional, default for every call
+  "samples":     [Sample, ...] }       required, non-empty
+
+Sample = { "calls":         [Call, ...],         required, non-empty, run in order
+           "tol":           [atol, rtol] }       optional, overrides the set's
+
+Call   = { "method":        str,                 required, a program method
+           "in":            [Value, ...],        required, positional arguments
+           "out":           [Value, ...],        optional, positional outputs
+           "intermediates": {name: Key, ...},    optional
+           "tol":           [atol, rtol] }       optional, overrides the sample's
+
+Value  = Key | number | true | false | [Value, ...]
+Key    = str naming a tensor in the set's data.safetensors
+```
+
+- Each sample starts from the program's initial state. Samples are independent and
+  may run in any order; the calls within a sample run in order with no reset
+  between them, so each sees the state the previous one left, whichever method
+  that call ran.
+- A tensor's key is its JSON path in the index, e.g. `samples[0].calls[1].in[2]`.
+  Below a call that is `.in[k]`, `.out[k]` or `.intermediates.<name>`, with
+  `[l]` appended for each level of list nesting. Every key resolves, and every
+  tensor is referenced exactly once.
+- A JSON integer is an int scalar, any other number a float.
+- Intermediate names are top-level SSA value names of the call's method
+  (`[A-Za-z_][A-Za-z0-9_]*`).
+  Producers record what they can; consumers compare the names both sides have.
+- Tolerances are finite and non-negative, and advisory: a consumer may override them.
+  `write_samples` warns when atol or rtol is above 1e-3.
+- Readers reject another major `version` and accept any minor one: a minor bump
+  only adds fields, and readers ignore unknown fields and the tensors only those
+  fields reference.
+
+Not in v1:
+
+- Tensors shared across sets.
+- Values inside nested graphs.
+- Quantized method inputs and outputs (a `QuantSpec` on graph I/O). Integer and
+  float I/O of quantized models is supported; torch quantized dtypes and tensor
+  subclasses are not.
+
+TODO: specify quantized I/O as the storage form the method binds (packed for
+sub-byte), with quant parameters read from the program rather than the sample.
+
 ## Runtime
 
 The C++ runtime reads a `.ptn` into an in-memory graph and hands each method to

backends/native/test/runtime/deserialize/test_package.cpp

@@ -17,6 +17,7 @@
 #include <stdexcept>
 #include <string>
 #include <string_view>
+#include <utility>
 #include <vector>
 
 #include <gtest/gtest.h>
@@ -56,7 +57,9 @@ class TempPackage {
   std::filesystem::path path_;
 
  public:
-  explicit TempPackage(std::vector<uint8_t> tensors = {}) {
+  explicit TempPackage(
+      std::vector<uint8_t> tensors = {},
+      const std::vector<std::pair<std::string, std::string>>& extra = {}) {
     path_ = std::filesystem::temp_directory_path() /
         ("ptn_package_" + std::to_string(reinterpret_cast<uintptr_t>(this)) +
          ".ptn");
@@ -74,6 +77,9 @@ class TempPackage {
     add(archive.get(), kProgramEntry, as_bytes(program_));
     add(archive.get(), kSafeTensorsEntry, ByteSpan(tensors));
     add(archive.get(), kAliasesEntry, as_bytes(aliases_));
+    for (const auto& [name, bytes] : extra) {
+      add(archive.get(), name.c_str(), as_bytes(bytes));
+    }
     if (zip_close(archive.get()) != 0) {
       throw std::runtime_error("failed to close test package");
     }
@@ -135,6 +141,24 @@ TEST(PackageTest, LoadsConstantsOnDemand) {
   EXPECT_NO_THROW(package.verify());
 }
 
+TEST(PackageTest, IgnoresSampleMembers) {
+  const std::string index =
+      R"({"version":"1.0","samples":[{"calls":[{"method":"forward","in":["samples[0].calls[0].in[0]"]}]}]})";
+  const std::vector<uint8_t> data = make_safetensors(
+      R"({"samples[0].calls[0].in[0]":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})",
+      "in");
+  const TempPackage file(
+      {},
+      {{"samples/golden/index.json", index},
+       {"samples/golden/data.safetensors",
+        std::string(data.begin(), data.end())}});
+  const Package package = Package::load(file.path());
+
+  EXPECT_EQ(package.owner_keys(), (std::vector<std::string>{"weight", "bias"}));
+  EXPECT_FALSE(package.constant_info("samples[0].calls[0].in[0]"));
+  EXPECT_NO_THROW(package.verify());
+}
+
 TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) {
   const TempPackage file;
   const Package package = Package::load(file.path());

exir/_serialize/BUCK

@@ -38,6 +38,21 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+# Optional recorded method calls inside a .ptn; see backends/native/README.md.
+fbcode_target(_kind = runtime.python_library,
+    name = "ptn_samples",
+    typing = True,
+    srcs = [
+        "_ptn_samples.py",
+    ],
+    visibility = ["PUBLIC"],
+    deps = [
+        ":ptn",
+        "fbsource//third-party/pypi/safetensors:safetensors",
+        "//caffe2:torch",
+    ],
+)
+
 # so we can have access to EXECUTORCH_CLIENTS list.
 fbcode_target(_kind = runtime.python_library,
     name = "lib",

exir/_serialize/_ptn.py

@@ -243,16 +243,20 @@ def _safetensors_header(owners: dict[str, torch.Tensor]) -> tuple[bytes, list[st
     return len(encoded).to_bytes(8, "little") + encoded, order
 
 
-def _write_safetensors(pkg: zipfile.ZipFile, owners: dict[str, torch.Tensor]) -> None:
-    """Stream the safetensors member into `pkg` with its data section aligned."""
+def _write_safetensors(
+    pkg: zipfile.ZipFile,
+    owners: dict[str, torch.Tensor],
+    entry: str = SAFETENSORS_ENTRY,
+) -> None:
+    """Stream a safetensors member into `pkg` with its data section aligned."""
     header, order = _safetensors_header(owners)
     size = len(header) + sum(t.numel() * t.element_size() for t in owners.values())
     zip64 = size > zipfile.ZIP64_LIMIT
 
-    zinfo = zipfile.ZipInfo(SAFETENSORS_ENTRY, time.localtime(time.time())[:6])
+    zinfo = zipfile.ZipInfo(entry, time.localtime(time.time())[:6])
     local_header_size = (
         _LOCAL_HEADER_SIZE
-        + len(SAFETENSORS_ENTRY.encode("utf-8"))
+        + len(entry.encode("utf-8"))
         + (_ZIP64_LOCAL_EXTRA_SIZE if zip64 else 0)
     )
     assert pkg.fp is not None