export_training_data.py

#!/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")