| 75 | ) |
| 76 | |
| 77 | def run(dir: str) -> bytes: |
| 78 | tosaname = "out.tosa" |
| 79 | tosa_path = os.path.join(dir, tosaname) |
| 80 | with open(tosa_path, "wb") as f: |
| 81 | f.write(tosa_flatbuffer) |
| 82 | |
| 83 | # invoke vela |
| 84 | output_dir = os.path.join(dir, "output") |
| 85 | args.append(f"--output-dir={output_dir}") |
| 86 | args.append(tosa_path) |
| 87 | if verbose: |
| 88 | args.append("--verbose-all") |
| 89 | vela.main(" ".join(args).split(" ")) |
| 90 | |
| 91 | np_path = os.path.join(dir, "output", "out_vela.npz") |
| 92 | |
| 93 | blocks = b"" |
| 94 | with np.load(np_path, allow_pickle=False) as data: |
| 95 | # Construct our modified output_blocks with data in a form easily |
| 96 | # digested on the device side |
| 97 | bin_blocks = {"vela_bin_stream": b""} |
| 98 | |
| 99 | # copy command data through unmodified |
| 100 | bin_blocks["cmd_data"] = data["cmd_data"].tobytes() |
| 101 | |
| 102 | # copy weight data through unmodified |
| 103 | bin_blocks["weight_data"] = data["weight_data"].tobytes() |
| 104 | |
| 105 | # Add a block for scratch, inputs and outputs; scratch shape is a 1 element |
| 106 | # array giving us size in bytes so extract this and add a block of 0's. |
| 107 | # Currently we preallocated this on the host to provide SRAM for computation. |
| 108 | if not isinstance(data["scratch_shape"][0], np.int64): |
| 109 | raise RuntimeError("Expected scratch to be int64") |
| 110 | block_length = int(data["scratch_shape"][0]) |
| 111 | bin_blocks["scratch_size"] = struct.pack("<I", block_length) |
| 112 | |
| 113 | # Capture inputs and outputs |
| 114 | bin_blocks["inputs"] = vela_bin_pack_io("input", data) |
| 115 | bin_blocks["outputs"] = vela_bin_pack_io("output", data) |
| 116 | |
| 117 | bin_blocks["vela_end_stream"] = b"" |
| 118 | |
| 119 | # Emit the NPZ regions as: |
| 120 | # - 16 byte block name null terminated string (padded to 16 if name shorter) |
| 121 | # - 4 bytes of int32 block length and 12 bytes of 0's |
| 122 | # - block data (padded to 16 byte alignment at end) |
| 123 | # Repeat for all blocks |
| 124 | for key in bin_blocks.keys(): |
| 125 | block_name = bytes(key, "utf8")[:15] |
| 126 | block_name = block_name + b"\x00" * (16 - len(block_name)) |
| 127 | |
| 128 | # We need the acual unpadded block lengths for hw setup |
| 129 | block_length_bytes = struct.pack("<iiii", len(bin_blocks[key]), 0, 0, 0) |
| 130 | |
| 131 | # Pad block data to multiple of 16 bytes |
| 132 | block_data = bin_blocks[key] |
| 133 | block_data = block_data + b"\x00" * (15 - (len(block_data) - 1) % 16) |
| 134 | |