ECE 408 sem project

kernels_op_7/sweep.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+import csv
+import re
+import subprocess
+import sys
+from itertools import product
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+BUILD = ROOT / "kernels_op_7" / "build"
+CSV_PATH = ROOT / "kernels_op_7" / "sweep.csv"
+
+# Just BR x BC. 9 configs. BR=128 will OOM on smem (48KB cap), expected NOT_FOUND.
+GRID = {
+    "br": [32, 64, 128],
+    "bc": [32, 64, 128],
+}
+
+# from Makefile
+NVCC = "nvcc"
+CFLAGS = "-O3 -arch=sm_86 -std=c++17 -rdc=true -g -lineinfo"
+INCLUDES = "-I/usr/local/cuda/include -I/work/hdd/bche/hshah6/cutlass/include"
+LIBS = "-L/usr/local/cuda/lib64 -lcublas -lcublasLt"
+
+KERNEL_NAME = "flash_attention_forward_kernel"
+
+
+def run(cmd):
+    return subprocess.run(cmd, shell=True, cwd=ROOT, capture_output=True, text=True)
+
+
+def compile_one(tag, defs):
+    """Two-step compile matching the Makefile."""
+    obj = BUILD / f"{tag}.o"
+    binary = BUILD / tag
+    r = run(f"{NVCC} {CFLAGS} {INCLUDES} {defs} -c test_gpt2.cu -o {obj}")
+    if r.returncode != 0:
+        return None, f"COMPILE FAILED:\n{r.stderr}"
+    r = run(f"{NVCC} {CFLAGS} -o {binary} {obj} {LIBS}")
+    if r.returncode != 0:
+        return None, f"LINK FAILED:\n{r.stderr}"
+    return binary, ""
+
+
+def profile_one(binary, tag):
+    report = BUILD / tag
+    log = BUILD / f"{tag}.log"
+    cmd = (
+        f"nsys profile --stats=true --force-overwrite=true "
+        f"--cuda-event-trace=false -o {report} {binary}"
+    )
+    r = run(cmd)
+    log.write_text((r.stdout or "") + "\n" + (r.stderr or ""))
+    return log
+
+
+def parse_avg_ns(log_path, kernel_name):
+    pattern = re.compile(rf"^\s*\d.*\b{re.escape(kernel_name)}\b")
+    for line in log_path.read_text().splitlines():
+        if pattern.search(line):
+            cols = line.split()
+            if len(cols) >= 4:
+                return cols[3]
+    return "NOT_FOUND"
+
+
+def main():
+    BUILD.mkdir(parents=True, exist_ok=True)
+    configs = list(product(GRID["br"], GRID["bc"]))
+
+    first_fail_printed = False
+    with CSV_PATH.open("w", newline="") as fp:
+        out = csv.writer(fp)
+        out.writerow(["BR", "BC", "fa2_avg_ns"])
+
+        for i, (br, bc) in enumerate(configs, 1):
+            tag = f"br{br}_bc{bc}"
+
+            defs = f"-DSWEEP_BR={br} -DSWEEP_BC={bc}"
+            binary, err = compile_one(tag, defs)
+
+            if binary is None:
+                if not first_fail_printed:
+                    print(err)
+                    first_fail_printed = True
+                out.writerow([br, bc, "BUILD_FAIL"])
+                fp.flush()
+                continue
+
+            log = profile_one(binary, tag)
+            out.writerow([br, bc, parse_avg_ns(log, KERNEL_NAME)])
+            fp.flush()
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())