even faster on pimcomp models

validation/raptor_validation/pimsim_nn.py

@@ -1,8 +1,11 @@
 import json
 import re
 import shutil
+import struct
 from pathlib import Path
 
+import numpy as np
+
 
 _METRIC_PATTERNS = {
     "output_count": r"output count:\s+([0-9]+)\s+samples",
@@ -14,6 +17,14 @@ _METRIC_PATTERNS = {
 }
 
 
+def read_raptor_instruction_count(path):
+    with open(path, "rb") as f:
+        magic, version, count = struct.unpack("<4sII", f.read(12))
+    if magic != b"PIMB" or version != 1:
+        raise ValueError(f"Invalid Raptor binary instruction file: {path}")
+    return count
+
+
 def parse_pimsim_nn_metrics(output):
     metrics = {"raw_output": output}
     for name, pattern in _METRIC_PATTERNS.items():
@@ -45,6 +56,46 @@ def export_raptor_latency_artifact(pim_dir, output_dir):
         json.dump(config, f, separators=(",", ":"))
         f.write("\n")
 
+    binary_sources = sorted(pim_dir.glob("core_*.pim"), key=lambda path: int(path.stem.split("_")[1]))
+    if binary_sources:
+        record_dtype = np.dtype([
+            ("opcode", "u1"), ("rd", "u1"), ("r1", "u1"), ("flags", "u1"),
+            ("r2_or_imm", "<i4"), ("generic1", "<i4"),
+            ("generic2", "<i4"), ("generic3", "<i4"),
+        ])
+        for source in binary_sources:
+            destination = output_dir / source.name
+            shutil.copyfile(source, destination)
+            count = read_raptor_instruction_count(destination)
+            if destination.stat().st_size != 12 + count * record_dtype.itemsize:
+                raise ValueError(f"Invalid Raptor binary instruction file: {source}")
+            records = np.memmap(destination, dtype=record_dtype, mode="r+", offset=12, shape=(count,))
+            scalar_sldi = np.zeros(count, dtype=bool)
+            vmv = records["opcode"] == 22
+            sldi = records["opcode"] == 1
+            for register in np.unique(records["r2_or_imm"][vmv]):
+                definitions = np.flatnonzero(sldi & (records["rd"] == register))
+                consumers = np.flatnonzero(vmv & (records["r2_or_imm"] == register))
+                reaching = np.searchsorted(definitions, consumers) - 1
+                scalar_sldi[definitions[reaching[reaching >= 0]]] = True
+            for start in range(0, count, 1_000_000):
+                chunk = records[start:start + 1_000_000]
+                setbw = chunk["opcode"] == 8
+                chunk["generic1"][setbw] = 8
+                chunk["generic2"][setbw] = 8
+                address_sldi = (chunk["opcode"] == 1) & ~scalar_sldi[start:start + len(chunk)]
+                if np.any(chunk["r2_or_imm"][address_sldi] % 4):
+                    raise ValueError(f"Raptor address is not aligned to its fp32 element width: {source}")
+                chunk["r2_or_imm"][address_sldi] //= 4
+                transfer = (chunk["opcode"] >= 25) & (chunk["opcode"] <= 30)
+                if np.any(chunk["generic2"][transfer] % 4) or np.any(chunk["generic3"][transfer] % 4):
+                    raise ValueError(f"Raptor transfer field is not aligned to its fp32 element width: {source}")
+                chunk["generic2"][transfer] //= 4
+                chunk["generic3"][transfer] //= 4
+            records.flush()
+            del records
+        return output_dir
+
     byte_size_fields = {
         "ld": "size",
         "st": "size",
@@ -56,12 +107,19 @@ def export_raptor_latency_artifact(pim_dir, output_dir):
     for source in sorted(pim_dir.glob("core_*.json"), key=lambda path: int(path.stem.split("_")[1])):
         with open(source, encoding="utf-8") as f:
             instructions = json.load(f)
-        for instruction in instructions:
+        last_sldi = {}
+        scalar_sldi = set()
+        for index, instruction in enumerate(instructions):
+            if instruction["op"] == "sldi":
+                last_sldi[instruction["rd"]] = index
+            elif instruction["op"] == "vmv" and instruction["rs2"] in last_sldi:
+                scalar_sldi.add(last_sldi[instruction["rs2"]])
+        for index, instruction in enumerate(instructions):
             op = instruction["op"]
             if op == "setbw":
                 instruction["ibiw"] = 8
                 instruction["obiw"] = 8
-            elif op == "sldi":
+            elif op == "sldi" and index not in scalar_sldi:
                 instruction["imm"] = int8_bytes(instruction["imm"], "address")
             if field := byte_size_fields.get(op):
                 instruction[field] = int8_bytes(instruction[field], f"{op} {field}")