# Makes karpathy/nanochat run on an Intel GPU with native PyTorch XPU (tested on an Intel UHD 770, Xe-LP),
# and fixes three nanochat issues that show up with a short context (max_seq_len=512, like runs/runcpu.sh).
# Usage: python3 nanochat_xpu_patch.py /path/to/nanochat   (idempotent; tested on commit 92d63d4)
# https://linuxgpu.eu/blog/  -  Linux GPU
import re, pathlib, sys
root = pathlib.Path(sys.argv[1])
def sub(path, old, new, count=1):
    p = root / path; s = p.read_text()
    if new in s: return  # already patched
    assert old in s, (path, old)
    p.write_text(s.replace(old, new, count))
# --- common.py
sub("nanochat/common.py", """    elif torch.backends.mps.is_available():
        device_type = "mps\"""", """    elif hasattr(torch, "xpu") and torch.xpu.is_available():
        device_type = "xpu"
    elif torch.backends.mps.is_available():
        device_type = "mps\"""")
sub("nanochat/common.py", 'assert device_type in ["cuda", "mps", "cpu"]', 'assert device_type in ["cuda", "xpu", "mps", "cpu"]')
sub("nanochat/common.py", """        torch.cuda.manual_seed(42)
""", """        torch.cuda.manual_seed(42)
    if device_type == "xpu":
        assert torch.xpu.is_available(), "PyTorch not built with XPU"
        torch.xpu.manual_seed(42)
""")
# --- base_train / chat_sft: sync, memory, name
for f in ["scripts/base_train.py", "scripts/chat_sft.py"]:
    sub(f, 'synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None',
           'synchronize = torch.cuda.synchronize if device_type == "cuda" else (torch.xpu.synchronize if device_type == "xpu" else (lambda: None))')
    sub(f, 'get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0',
           'get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else (torch.xpu.max_memory_allocated if device_type == "xpu" else (lambda: 0))')
    sub(f, 'if device_type == "cuda":\n    gpu_device_name = torch.cuda.get_device_name(0)',
           'if device_type == "xpu":\n    print0(f"XPU device: {torch.xpu.get_device_name(0)}")\nif device_type == "cuda":\n    gpu_device_name = torch.cuda.get_device_name(0)')
for f in ["scripts/chat_cli.py", "scripts/chat_eval.py"]:
    sub(f, "choices=['cuda', 'cpu', 'mps']", "choices=['cuda', 'xpu', 'cpu', 'mps']")
# --- gpt.py : scalar ** tensor => kernel fp64, absent on Xe-LP (UHD 770) -> UR error 45
sub("nanochat/gpt.py", "inv_freq = 1.0 / (base ** (channel_range / head_dim))",
    "inv_freq = torch.exp(channel_range * (-math.log(base) / head_dim))  # xpu: no fp64 on Xe-LP")
g = root / "nanochat/gpt.py"; gs = g.read_text()
if "\nimport math\n" not in gs:
    g.write_text(gs.replace("\nimport torch\n", "\nimport math\nimport torch\n", 1))
# --- engine.py : topk over 32k vocab exceeds max work-group (512) on Xe-LP -> do topk on CPU
sub("nanochat/engine.py", "        vals, idx = torch.topk(logits, k, dim=-1)\n",
    "        if logits.device.type == \"xpu\":  # Xe-LP: work-group max 512 -> topk on CPU\n"
    "            vals, idx = torch.topk(logits.float().cpu(), k, dim=-1)\n"
    "            vals, idx = vals.to(logits.device), idx.to(logits.device)\n"
    "        else:\n"
    "            vals, idx = torch.topk(logits, k, dim=-1)\n")
# --- chat_sft.py: a fully masked micro-batch gives mean CE = 0/0 = nan, and the smoothed loss stays nan forever
#     (not XPU-specific, reproduced on CPU). Use sum / max(n_valid, 1).
sub("scripts/chat_sft.py", """        loss = model(x, y)
        train_loss = loss.detach() # for logging
""", """        n_valid = (y >= 0).sum()  # fully masked micro-batch -> mean CE = 0/0 = nan: use sum / max(n,1) (= mean, or 0 with zero grads)
        loss = model(x, y, loss_reduction='sum') / n_valid.clamp(min=1)
        if n_valid.item() > 0:
            train_loss = loss.detach() # for logging (keep previous value on fully masked batches)
""")
# --- chat_sft.py: skip updates with a non-finite gradient; NANOCHAT_CLIP=<max norm> clips the gradient,
#     NANOCHAT_LOG_GN=1 logs the gradient norm and raw loss at every step
sub("scripts/chat_sft.py", """    # step the optimizer
""", """    # guard: skip non-finite gradients (one bad update otherwise turns every weight into nan)
    _gn = torch.nn.utils.clip_grad_norm_(orig_model.parameters(), max_norm=float(os.environ.get("NANOCHAT_CLIP", "1e9")))
    if not torch.isfinite(_gn):
        print0(f"WARNING step {step}: non-finite grad norm ({_gn.item()}), update skipped")
        for _p in orig_model.parameters():
            if _p.grad is not None:
                _p.grad.zero_()
    elif os.environ.get("NANOCHAT_LOG_GN"):
        print0(f"gradnorm step {step} gn {_gn.item():.4f} rawloss {train_loss.item():.4f}")
    # step the optimizer
""")
# --- chat_sft.py: nanochat bug with a short max_seq_len (not XPU-specific): the best-fit loader only takes
#     conversations that fit WHOLE in 513 tokens; longer ones stay in the 100-entry buffer until it is full of them,
#     then every row is padding (zero target tokens, gradient norm 0 from step ~6). Muon keeps pushing its stale
#     momentum, the model drifts, then goes nan. Fix: crop each conversation to row_capacity.
sub("scripts/chat_sft.py", """            ids, mask = tokenizer.render_conversation(conversation)
            conv_buffer.append((ids, mask))
""", """            ids, mask = tokenizer.render_conversation(conversation)
            ids, mask = ids[:row_capacity], mask[:row_capacity]  # crop: with a short max_seq_len, longer convs never fit and clog the buffer
            conv_buffer.append((ids, mask))
""")
print("patched")