#!/usr/bin/env python3 """PAX Training Data Export Pipeline — Lean 4 + PTX + Futhark + Spec → JSONL""" import json, re, hashlib, random, os from pathlib import Path ROOT = Path(__file__).parent PROMPT_TEMPLATES = { "fp16": [ "Write a Lean 4 formalization of IEEE-754 binary16 RNE with proven |round(x)-x| ≤ 0.5 ulp for FP16 GEMM on Ampere sm_86.", "Implement FP16 addition, multiplication, and FMA in Lean 4 with proven rounding error bounds matching __hadd, __hmul, __hfma.", "Formalize FP16→FP32 exact conversion for GEMM accumulation. Prove toRat(toFloat32(x)) = toRat(x) for all normal/subnormal FP16.", ], "gemm": [ "Write a verified 128×128 GEMM kernel for RTX 3080 sm_86 using mma.sync.aligned.m16n8k8 FP16→FP32.", "Prove PTX mma.sync semantics match WMMA abstract machine. Register-level equivalence for FP16→FP32.", "Write a PAX-compliant GEMM with ldmatrix.x4, mma.sync, cp.async and Lean 4 PO1+PO3+PO5 proofs.", ], "pipeline": [ "Define a 3-stage async cp.async pipeline in Lean 4 with proven throughput bound ≥ (1-1/stages)×min(bw_compute, bw_memory).", "Formalize cp.async.ca.shared.global with commit/wait_all. Prove pipeline preserves happens-before ordering across stages.", "Write a 3-stage async GEMM pipeline for RTX 3080 with proven overlap bound and Lean 4 PO4+PO6+PO7 certificates.", ], "epilogue": [ "Define epilogue fusion algebra: Fuse(BiasAdd, GeLU) ≡ GeLU ∘ BiasAdd. Prove |GeLU_approx - GeLU_exact| ≤ 0.001.", "Formalize in-register Bias+GeLU and Residual+GeLU fusion. Prove register bound: regs(fuse) ≤ regs(f) + regs(g) + 8.", "Write an Ampere epilogue kernel fusing BiasAdd+GeLU in a single pass with Lean 4 PO8 correctness certificate.", ], "warp": [ "Write warp-level reduction using shfl.sync.xor.b32. Prove correctness for dot product and softmax.", "Formalize SIMT divergence and reconvergence stack. Prove warp reconverges before barrier.", ], "architecture": [ "Map PAX Architecture axioms to Lean 4 proof obligations: Axiom 1→PO1, Axiom 2→PO2, Axiom 3→PO3, Axiom 4→PO4, Axiom 5→PO5+PO8.", "Explain the HyperKitty Constraint DAG and its Lean 4 formalization in PAX/ConstraintDAG.lean.", ], } CONSTRAINTS = { "fp16": ["PO4", "PO5"], "gemm": ["PO1", "PO3", "PO5", "PO8"], "pipeline": ["PO4", "PO6", "PO7"], "epilogue": ["PO8"], "index_space": ["PO1", "PO2"], "warp": ["PO3", "PO4"], "architecture":["PO8"], } SOURCE_FILES = [ ("PAX/ConstraintDAG.lean", "architecture", "all"), ("PAX/IR_DAG.lean", "architecture", "all"), ("PAX/PipelineDAG.lean", "pipeline", "sm_86"), ("PAX/Float16_Rounding.lean", "fp16", "sm_86"), ("PAX/WMMA.lean", "gemm", "sm_86"), ("PAX/TrainingData.lean", "architecture", "all"), ("src/rtx_gemm_ptx.cu", "gemm", "sm_86"), ("src/rtx_gemm_pipeline.cu", "pipeline", "sm_86"), ("src/rtx_gemm_epilogue.cu", "epilogue", "sm_86"), ("src/pax_kernel.fut", "gemm", "sm_86"), ("docs/PAX_ARCHITECTURE.md", "architecture", "all"), ] def extract_lean_theorems(content): pattern = r'(theorem|lemma)\s+(\w+)([^:=]*:[^:=]*):=\s*(by[^\n]*(?:\n [^\n]*)*)' return re.findall(pattern, content, re.MULTILINE) def extract_ptx_kernels(content): pattern = r'(__global__[^\{]*\{[^\}]*\})' return re.findall(pattern, content, re.DOTALL) def make_id(s): return hashlib.md5(s.encode()).hexdigest()[:12] def generate_examples(): examples = [] for rel_path, category, arch in SOURCE_FILES: path = ROOT / rel_path if not path.exists(): continue content = path.read_text(encoding="utf-8", errors="replace") prompts = PROMPT_TEMPLATES.get(category, PROMPT_TEMPLATES["architecture"]) if rel_path.endswith(".lean"): for kind, name, sig, proof in extract_lean_theorems(content): thm = f"{kind} {name}{sig}" for prompt in prompts[:2]: examples.append({ "id": make_id(rel_path + name), "instruction": prompt, "input": f"Arch: {arch} | Category: {category} | Constraints: {' '.join(CONSTRAINTS.get(category, []))}", "output": f"```lean4\n{thm} := {proof}\n```", "metadata": {"file": rel_path, "arch": arch, "category": category, "constraints": CONSTRAINTS.get(category, [])}, }) elif rel_path.endswith(".cu"): kernels = extract_ptx_kernels(content) for kernel in kernels: for prompt in prompts[:2]: examples.append({ "id": make_id(rel_path + kernel[:40]), "instruction": prompt, "input": f"Arch: {arch} | Category: {category} | Constraints: {' '.join(CONSTRAINTS.get(category, []))}", "output": f"```cuda\n{kernel[:2000]}\n```", "metadata": {"file": rel_path, "arch": arch, "category": category, "constraints": CONSTRAINTS.get(category, [])}, }) elif rel_path.endswith(".fut"): for prompt in prompts[:2]: examples.append({ "id": make_id(rel_path), "instruction": prompt, "input": f"Arch: {arch} | Category: {category} | Constraints: {' '.join(CONSTRAINTS.get(category, []))}", "output": f"```futhark\n{content[:2000]}\n```", "metadata": {"file": rel_path, "arch": arch, "category": category, "constraints": CONSTRAINTS.get(category, [])}, }) elif rel_path.endswith(".md"): sections = re.split(r'\n## ', content) for section in sections[:5]: title = section.split('\n')[0].strip("# ") for prompt in prompts[:1]: examples.append({ "id": make_id(rel_path + title), "instruction": prompt, "input": f"Arch: {arch} | Category: {category}", "output": f"```markdown\n{section[:1500]}\n```", "metadata": {"file": rel_path, "arch": arch, "category": category, "constraints": []}, }) # Dedup seen = set() unique = [] for ex in examples: key = ex["id"] if key not in seen: seen.add(key) unique.append(ex) return unique def split_and_write(examples): os.makedirs(ROOT / "build", exist_ok=True) random.seed(42) random.shuffle(examples) n = len(examples) splits = { "train": examples[:int(0.90 * n)], "val": examples[int(0.90 * n):int(0.95 * n)], "test": examples[int(0.95 * n):], } for name, data in splits.items(): out = ROOT / "build" / f"pax_{name}.jsonl" with open(out, "w", encoding="utf-8") as f: for ex in data: f.write(json.dumps(ex) + "\n") print(f" {name}: {len(data)} examples -> {out}") if __name__ == "__main__": print("=== PAX Training Data Extraction ===") examples = generate_examples() print(f"Total unique examples: {len(examples)}") split_and_write(examples) print("Done. Run: python3 finetune_pax_coder.py")