165 lines
6.0 KiB
Python
165 lines
6.0 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import random
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
|
|
def u32(value):
|
|
return value & 0xFFFFFFFF
|
|
|
|
|
|
def split_u66(value):
|
|
return [
|
|
value & 0xFFFFFFFF,
|
|
(value >> 32) & 0xFFFFFFFF,
|
|
(value >> 64) & 0x3,
|
|
]
|
|
|
|
|
|
def matmul4x4(a, b):
|
|
c = []
|
|
for row in range(4):
|
|
for col in range(4):
|
|
acc = 0
|
|
for k in range(4):
|
|
acc += a[row * 4 + k] * b[k * 4 + col]
|
|
c.append(acc)
|
|
return c
|
|
|
|
|
|
def write_words_bin(path, words):
|
|
with path.open("wb") as f:
|
|
for word in words:
|
|
f.write(struct.pack("<I", u32(word)))
|
|
|
|
|
|
def write_words_hex(path, words):
|
|
with path.open("w", encoding="utf-8") as f:
|
|
for word in words:
|
|
f.write(f"{u32(word):08x}\n")
|
|
|
|
|
|
def build_vectors(groups, rng, max_value, keep_vectors, include_zero_matrix):
|
|
vectors = []
|
|
source_words = []
|
|
result_words = []
|
|
|
|
for group in range(groups):
|
|
if include_zero_matrix and group == 0:
|
|
a = [0] * 16
|
|
b = [0] * 16
|
|
else:
|
|
a = [rng.randrange(max_value + 1) for _ in range(16)]
|
|
b = [rng.randrange(max_value + 1) for _ in range(16)]
|
|
c = matmul4x4(a, b)
|
|
|
|
source_words.extend(a)
|
|
source_words.extend(b)
|
|
for value in c:
|
|
result_words.extend(split_u66(value))
|
|
|
|
if keep_vectors:
|
|
vectors.append({
|
|
"group": group,
|
|
"A": a,
|
|
"B": b,
|
|
"C_u66": [str(value) for value in c],
|
|
"C_words": [split_u66(value) for value in c],
|
|
})
|
|
|
|
return vectors, source_words, result_words
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate full 4MB 32-bit ExtRAM images for the matrix-multiply task."
|
|
)
|
|
parser.add_argument("--groups", type=int, default=10,
|
|
help="number of A/B matrix groups")
|
|
parser.add_argument("--seed", type=int, default=1, help="random seed")
|
|
parser.add_argument("--max-value", type=lambda x: int(x, 0), default=0xFFFFFFFF,
|
|
help="maximum unsigned 32-bit input value, inclusive")
|
|
parser.add_argument("--no-zero-matrix", action="store_true",
|
|
help="do not force group0 A/B to be all-zero matrices")
|
|
parser.add_argument("--memory-bytes", type=lambda x: int(x, 0), default=4 * 1024 * 1024,
|
|
help="ExtRAM image size in bytes")
|
|
parser.add_argument("--out-dir", type=Path, default=Path("matmul_testdata"),
|
|
help="output directory")
|
|
parser.add_argument("--dump-vectors", action="store_true",
|
|
help="include every A/B/C matrix in matmul_vectors.json")
|
|
args = parser.parse_args()
|
|
|
|
if args.memory_bytes <= 0 or args.memory_bytes % 4 != 0:
|
|
raise SystemExit("--memory-bytes must be a positive multiple of 4")
|
|
memory_words = args.memory_bytes // 4
|
|
max_groups = memory_words // 80
|
|
|
|
if args.groups <= 0:
|
|
raise SystemExit("--groups must be positive")
|
|
if args.groups > max_groups:
|
|
raise SystemExit(f"--groups={args.groups} does not fit in {args.memory_bytes} bytes; max is {max_groups}")
|
|
if args.max_value < 0 or args.max_value > 0xFFFFFFFF:
|
|
raise SystemExit("--max-value must be in range 0..0xffffffff")
|
|
|
|
rng = random.Random(args.seed)
|
|
vectors, source_words, result_words = build_vectors(
|
|
args.groups, rng, args.max_value, args.dump_vectors, not args.no_zero_matrix
|
|
)
|
|
dst_base_word = len(source_words)
|
|
result_end_word = dst_base_word + len(result_words)
|
|
|
|
input_image_words = [0] * memory_words
|
|
expected_image_words = [0] * memory_words
|
|
input_image_words[0:len(source_words)] = source_words
|
|
expected_image_words[0:len(source_words)] = source_words
|
|
expected_image_words[dst_base_word:result_end_word] = result_words
|
|
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
write_words_bin(args.out_dir / "matmul_source_ab.bin", source_words)
|
|
write_words_bin(args.out_dir / "matmul_expected_c.bin", result_words)
|
|
write_words_bin(args.out_dir / "extram_input_4mb.bin", input_image_words)
|
|
write_words_bin(args.out_dir / "extram_expected_4mb.bin", expected_image_words)
|
|
write_words_hex(args.out_dir / "matmul_source_ab.hex", source_words)
|
|
write_words_hex(args.out_dir / "matmul_expected_c.hex", result_words)
|
|
write_words_hex(args.out_dir / "extram_input_4mb.hex", input_image_words)
|
|
write_words_hex(args.out_dir / "extram_expected_4mb.hex", expected_image_words)
|
|
|
|
metadata = {
|
|
"groups": args.groups,
|
|
"seed": args.seed,
|
|
"max_value": args.max_value,
|
|
"zero_matrix_group": None if args.no_zero_matrix else 0,
|
|
"memory_bytes": args.memory_bytes,
|
|
"memory_words": memory_words,
|
|
"max_groups_for_memory": max_groups,
|
|
"source_words": len(source_words),
|
|
"result_words": len(result_words),
|
|
"src_base_offset": 0,
|
|
"dst_base_offset": dst_base_word * 4,
|
|
"unused_tail_words": memory_words - result_end_word,
|
|
"bytes_per_source_group": 128,
|
|
"bytes_per_result_group": 192,
|
|
"input_image": "source region filled, result region and unused tail filled with 0",
|
|
"expected_image": "source region preserved, result region filled, unused tail filled with 0",
|
|
"layout": "A0,B0,A1,B1,...,A(n-1),B(n-1),C0,C1,...,C(n-1),zero padding",
|
|
}
|
|
if args.dump_vectors:
|
|
metadata["vectors"] = vectors
|
|
(args.out_dir / "matmul_vectors.json").write_text(
|
|
json.dumps(metadata, indent=2), encoding="utf-8"
|
|
)
|
|
|
|
print(f"generated {args.groups} groups in {args.out_dir}")
|
|
print(f"memory size: {args.memory_bytes} bytes ({memory_words} words)")
|
|
print("SRC_BASE offset: 0")
|
|
print(f"DST_BASE offset: {dst_base_word * 4}")
|
|
print(f"unused tail: {memory_words - result_end_word} words")
|
|
print("input image: extram_input_4mb.bin / extram_input_4mb.hex")
|
|
print("expected image: extram_expected_4mb.bin / extram_expected_4mb.hex")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|