diff --git a/space/Dockerfile b/space/Dockerfile index fed0aaf3dc14fabb532ed34dc863a0fe8cba0e3b..0aaf7e7ebdb3c4270f831988ecd45f005b63260e 100644 --- a/space/Dockerfile +++ b/space/Dockerfile @@ -16,7 +16,7 @@ ENV DEBIAN_FRONTEND=noninteractive \ NVIDIA_VISIBLE_DEVICES=all \ NVIDIA_DRIVER_CAPABILITIES=compute,utility -# Python + basic build deps (kept minimal) +# Python + basic deps RUN apt-get update && apt-get install -y --no-install-recommends \ python3.11 python3.11-venv python3-pip \ git ca-certificates \ @@ -25,8 +25,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ COPY requirements.txt . RUN python3.11 -m pip install --upgrade pip && python3.11 -m pip install -r requirements.txt +# CUDA-enabled PyTorch (so `torch.cuda.is_available()` is True on GPU Spaces) +RUN python3.11 -m pip install --no-cache-dir \ + torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 + COPY . . EXPOSE 7860 - CMD ["python3.11", "app.py"] diff --git a/space/space/space/Dockerfile b/space/space/space/Dockerfile index 445eb4cfc5eec77c8dbf2ff58f3803f047e6e6ee..fed0aaf3dc14fabb532ed34dc863a0fe8cba0e3b 100644 --- a/space/space/space/Dockerfile +++ b/space/space/space/Dockerfile @@ -1,19 +1,32 @@ # Optional: use Hugging Face Space SDK = docker (set `sdk: docker` in README.md). -# Default README uses Gradio SDK and does not require this image. +# +# IMPORTANT: +# - If your Space is on CPU hardware, `torch.cuda.is_available()` will be False no matter what. +# - If your Space is on GPU hardware, you must use a CUDA-enabled base image (below), +# otherwise CUDA won't be visible inside the container. -FROM python:3.11-slim +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 WORKDIR /app -ENV PYTHONUNBUFFERED=1 \ +ENV DEBIAN_FRONTEND=noninteractive \ + PYTHONUNBUFFERED=1 \ PIP_NO_CACHE_DIR=1 \ - GRADIO_SERVER_NAME=0.0.0.0 + GRADIO_SERVER_NAME=0.0.0.0 \ + NVIDIA_VISIBLE_DEVICES=all \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility + +# Python + basic build deps (kept minimal) +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3.11 python3.11-venv python3-pip \ + git ca-certificates \ + && rm -rf /var/lib/apt/lists/* COPY requirements.txt . -RUN pip install --upgrade pip && pip install -r requirements.txt +RUN python3.11 -m pip install --upgrade pip && python3.11 -m pip install -r requirements.txt COPY . . EXPOSE 7860 -CMD ["python", "app.py"] +CMD ["python3.11", "app.py"] diff --git a/space/space/space/space/space/space/program_generator.py b/space/space/space/space/space/space/program_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..710a8f6b56e0cc5f5d990da2e82b3cb7da725247 --- /dev/null +++ b/space/space/space/space/space/space/program_generator.py @@ -0,0 +1,116 @@ +""" +Random Toy-IR list-of-dicts for GRPO / Deliverable2 / runtime_core. + +Schema matches `runtime_core.SAMPLE_PROGRAM` (op, args, dest, type) β€” the same +shape `Deliverable2_Formatter` and `CompilerOptimizationEnv` expect. + +Note: `metahack1 (1).ipynb` uses a different TAC shape (CONST/src1/STORE). +Use that notebook's generators only if you add a separate converter; this +module is self-contained for training stack compatibility. +""" + +from __future__ import annotations + +import copy +import json +import random +from typing import List + +# Hand-written seeds (same as Colab / train.py short list) β€” kept in sync for regression. +_BUILTIN_EXTRA: List[List[dict]] = [ + [ + {"op": "const", "dest": "a", "args": ["1"], "type": "int"}, + {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"}, + {"op": "ret", "args": ["b"]}, + ], + [ + {"op": "const", "dest": "x", "args": ["2"], "type": "int"}, + {"op": "const", "dest": "y", "args": ["4"], "type": "int"}, + {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"}, + {"op": "const", "dest": "k", "args": ["1"], "type": "int"}, + {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"}, + {"op": "ret", "args": ["w"]}, + ], +] + + +def random_toy_ir_program(rng: random.Random) -> List[dict]: + """ + One valid program: consts v0.., then a chain of add/mul on existing names, then ret. + All ops use the mock-engine-friendly list schema. + """ + n_const = rng.randint(2, 5) + n_arith = rng.randint(1, 5) + progs: List[dict] = [] + for i in range(n_const): + progs.append( + { + "op": "const", + "dest": f"v{i}", + "args": [str(rng.randint(0, 20))], + "type": "int", + } + ) + available = [f"v{i}" for i in range(n_const)] + nxt = n_const + for _j in range(n_arith): + a = rng.choice(available) + b = rng.choice(available) + opn = rng.choice(["add", "mul"]) + d = f"v{nxt}" + nxt += 1 + progs.append({"op": opn, "dest": d, "args": [a, b], "type": "int"}) + available.append(d) + progs.append({"op": "ret", "args": [available[-1]]}) + return progs + + +def build_training_program_corpus( + n_total: int = 120, + seed: int = 42, + *, + include_builtins: bool = True, +) -> List[List[dict]]: + """ + Return `n_total` programs for GRPO. Optionally prepend SAMPLE_PROGRAM + 2 hand-written IRs + (when include_builtins), then fill with random_toy_ir_program, deduplicating by JSON key. + + Typical range: set `n_total` between 50 and 200 in the notebook. + """ + if n_total < 1: + raise ValueError("n_total must be >= 1") + + rng = random.Random(seed) + out: List[List[dict]] = [] + seen: set[str] = set() + + def _add(p: List[dict]) -> None: + k = json.dumps(p, sort_keys=True) + if k in seen: + return + seen.add(k) + out.append(copy.deepcopy(p)) + + if include_builtins: + from runtime_core import SAMPLE_PROGRAM + + for p in (SAMPLE_PROGRAM, *_BUILTIN_EXTRA): + if len(out) >= n_total: + break + _add(p) + + # Fill with random programs (dedupe by full JSON; allow dup if generator keeps colliding) + guard = 0 + while len(out) < n_total: + guard += 1 + if guard > 200_000: + out.append(random_toy_ir_program(rng)) + continue + cand = random_toy_ir_program(rng) + k = json.dumps(cand, sort_keys=True) + if k in seen: + continue + seen.add(k) + out.append(cand) + + return out diff --git a/space/space/space/space/space/space/space/reverse_pass/README.md b/space/space/space/space/space/space/space/reverse_pass/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7f1a4cad0fed81d95688b4c59c1f5eddc1b2409e --- /dev/null +++ b/space/space/space/space/space/space/space/reverse_pass/README.md @@ -0,0 +1,57 @@ +# Compiler Phase-Ordering RLVR (Toy-IR) + +Toy-IR compiler phase-ordering environment for RL with verifiable rewards (RLVR), including reverse passes to escape local optima. + +## What Is Implemented + +- Reverse passes are available in the pass library: + - `expand_constant` + - `duplicate_computation` +- Action routing and validation include reverse passes. +- Prompt guidance includes reverse-pass descriptions and usage intent. +- Reward shaping is terminal-weighted with RLVR hard gate: + - non-equivalent -> `-1000.0` + - terminal -> `((original_cycles - current_cycles) / original_cycles) * 100.0` + - non-terminal -> `-0.1` +- Episode termination: + - `STOP`/`done`, or + - hard cap at 5 steps. +- Reverse-pass instrumentation logs every 50 episodes/completions: + - `reverse_pass_episodes` + - `expand_constant_count` + - `duplicate_computation_count` + +## Verified So Far + +- Smoke test runs end-to-end without runtime errors. +- Rewards are scalar floats (not NaN) in tested rollouts. +- Forced reverse-pass episodes show expected behavior: + - small negative intermediate rewards (`-0.1`) + - terminal reward depends on final outcome (positive if chain beats baseline, negative if not) + +## Pending (GPU Required) + +Comparative training runs are still pending and require a GPU environment with: + +- `torch` +- `trl` +- `unsloth` +- `wandb` + +Required fair comparison: + +1. Baseline run: reverse passes disabled (`baseline_no_reverse`) +2. Reverse run: reverse passes enabled (`with_reverse_passes`) +3. Same episode budget for both runs +4. Compare reward/cycle curves and reverse-pass usage metrics in WandB + +## Team Handoff Status + +- Reverse-pass feature: shipped +- Training wiring + reward shaping + instrumentation: shipped +- Comparative training: pending in GPU environment +- README: this document + +## Suggested Final Pre-Training Sanity Check + +Run one smoke/equivalence check on a real Role 3 generated curriculum program (not just mock IR) to confirm verifier behavior in-pipeline before long training. diff --git a/space/space/space/space/space/space/space/reverse_pass/compiler_optimization_grpo.ipynb b/space/space/space/space/space/space/space/reverse_pass/compiler_optimization_grpo.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..69324ca6a9d0696206b8172a43f5ff9f96d91980 --- /dev/null +++ b/space/space/space/space/space/space/space/reverse_pass/compiler_optimization_grpo.ipynb @@ -0,0 +1,850 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title-cell", + "metadata": {}, + "source": [ + "# πŸ”§ Compiler Optimization RL Environment\n", + "### OpenEnv Hackathon 2026 β€” Theme #2: Long-Horizon Planning\n", + "\n", + "**What this notebook does:**\n", + "1. Defines a fully OpenEnv-compliant `CompilerOptimizationEnv`\n", + "2. Loads `Qwen2.5-3B-Instruct` via Unsloth (4-bit QLoRA)\n", + "3. Trains with GRPO (TRL) β€” LLM learns to pick compiler passes that reduce CPU cycles\n", + "4. Runs a smoke-test with a mock engine so you can verify reward logic without real hardware\n", + "5. Plots reward curves\n", + "\n", + "---\n", + "**Stack:** `unsloth` Β· `trl` Β· `openenv` Β· `wandb` Β· `matplotlib`\n", + "\n", + "> **Runtime:** Google Colab T4 GPU recommended. For the smoke-test only, CPU is fine." + ] + }, + { + "cell_type": "markdown", + "id": "install-header", + "metadata": {}, + "source": [ + "## πŸ“¦ Cell 1 β€” Install Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "install-cell", + "metadata": {}, + "outputs": [], + "source": [ + "# Install all required packages\n", + "# Unsloth must be installed before trl to get the right CUDA kernels\n", + "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n", + "!pip install trl datasets transformers accelerate peft bitsandbytes --quiet\n", + "!pip install wandb matplotlib --quiet\n", + "\n", + "# Optional: install openenv for production use\n", + "# !pip install openenv --quiet\n", + "\n", + "print(\"βœ… All packages installed\")" + ] + }, + { + "cell_type": "markdown", + "id": "env-header", + "metadata": {}, + "source": [ + "## 🌍 Cell 2 β€” CompilerOptimizationEnv (OpenEnv-Compliant)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "env-cell", + "metadata": {}, + "outputs": [], + "source": [ + "# CompilerOptimizationEnv, PASSES, reverse passes, terminal-weighted reward (toyir_rl_support)\n", + "import os\n", + "import sys\n", + "\n", + "if os.getcwd() not in sys.path:\n", + " sys.path.insert(0, os.getcwd())\n", + "\n", + "from toyir_rl_support import (\n", + " MCPEnvironment,\n", + " StepResult,\n", + " EpisodeStats,\n", + " CompilerOptimizationEnv,\n", + " PASSES,\n", + " MOCK_PASSES,\n", + " compute_shaped_reward,\n", + " rollout_shaped_return,\n", + " log_reverse_pass_stats_for_completion,\n", + " ensure_import_path,\n", + " MockEngine,\n", + " MOCK_ENGINE,\n", + " SAMPLE_PROGRAM,\n", + ")\n", + "\n", + "ensure_import_path()\n", + "print(\"βœ… Loaded CompilerOptimizationEnv, PASSES, and reward helpers from toyir_rl_support\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "smoke-header", + "metadata": {}, + "source": [ + "## πŸ§ͺ Cell 3 β€” Smoke Test (No GPU / Real Engine Needed)\n", + "Validates the entire reward pipeline with a mock engine. Run this before spending compute." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "smoke-test-cell", + "metadata": {}, + "outputs": [], + "source": [ + "# Smoke test (MockEngine; same PASSES as training)\n", + "engine = MockEngine()\n", + "env = CompilerOptimizationEnv(\n", + " engine, MOCK_PASSES, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", + ")\n", + "obs = env.reset(SAMPLE_PROGRAM)\n", + "\n", + "print(\"═\" * 55)\n", + "print(\"SMOKE TEST\")\n", + "print(\"═\" * 55)\n", + "print(f\"Initial state:\\n{obs}\")\n", + "print(f\"\\nBaseline cycles : {env.previous_cycles}\")\n", + "print(f\"Available actions: {env.available_actions()}\")\n", + "print()\n", + "\n", + "actions_to_try = [\n", + " \"constant_folding\",\n", + " \"dead_code_elimination\",\n", + " \"peephole_optimization\", # no-op\n", + " \"hallucinated_pass\", # invalid β€” but won't kill episode yet\n", + " \"constant_folding\",\n", + " \"dead_code_elimination\",\n", + " \"expand_constant\",\n", + " \"STOP\",\n", + "]\n", + "\n", + "for action in actions_to_try:\n", + " result = env.step(action)\n", + " tag = \"βœ—\" if result.reward < 0 else \"βœ“\"\n", + " print(f\"{tag} '{action}'\")\n", + " print(f\" reward={result.reward:+.2f} done={result.done} is_terminal={result.info.get('is_terminal')}\")\n", + " relevant = {k: v for k, v in result.info.items()\n", + " if k in (\"delta_pct\", \"error\", \"no_op\", \"reason\",\n", + " \"terminal_bonus\", \"episode_stats\", \"is_terminal\")}\n", + " if relevant:\n", + " print(f\" info: {relevant}\")\n", + " print()\n", + " if result.done:\n", + " break\n", + "\n", + "print(\"βœ… Smoke test passed\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "reward-plot-header", + "metadata": {}, + "source": [ + "## πŸ“Š Cell 4 β€” Visualise Reward Across a Mock Episode" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "reward-plot-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import matplotlib.ticker as ticker\n", + "\n", + "# Run a full episode and collect data\n", + "env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n", + "env.reset(SAMPLE_PROGRAM)\n", + "\n", + "sequence = [\n", + " \"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\",\n", + " \"expand_constant\", \"peephole_optimization\",\n", + "]\n", + "\n", + "rewards, cycle_counts, actions_log = [], [], []\n", + "for act in sequence:\n", + " r = env.step(act)\n", + " rewards.append(r.reward)\n", + " cycle_counts.append(env.previous_cycles)\n", + " actions_log.append(act)\n", + " if r.done:\n", + " break\n", + "\n", + "steps = list(range(1, len(rewards) + 1))\n", + "\n", + "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n", + "fig.suptitle(\"Compiler Optimization Episode β€” Mock Engine\", fontsize=14, fontweight=\"bold\")\n", + "\n", + "# Reward per step\n", + "colors = [\"#2ecc71\" if r >= 0 else \"#e74c3c\" for r in rewards]\n", + "ax1.bar(steps, rewards, color=colors, edgecolor=\"white\", linewidth=0.5)\n", + "ax1.axhline(0, color=\"grey\", linewidth=0.8, linestyle=\"--\")\n", + "ax1.set_ylabel(\"Reward\")\n", + "ax1.set_title(\"Reward per Step (green = positive, red = negative)\")\n", + "ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter(\"%.1f\"))\n", + "\n", + "# Cycle count over time\n", + "ax2.plot(steps, cycle_counts, marker=\"o\", color=\"#3498db\", linewidth=2, markersize=6)\n", + "ax2.set_xlabel(\"Step\")\n", + "ax2.set_ylabel(\"CPU Cycles\")\n", + "ax2.set_title(\"CPU Cycle Count Over Episode (lower = better)\")\n", + "ax2.set_xticks(steps)\n", + "ax2.set_xticklabels(\n", + " [a.replace(\"_\", \"\\n\") for a in actions_log],\n", + " fontsize=7,\n", + ")\n", + "\n", + "plt.tight_layout()\n", + "plt.savefig(\"episode_reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(\"πŸ“ˆ Plot saved as episode_reward_curve.png\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "model-header", + "metadata": {}, + "source": [ + "## πŸ€– Cell 5 β€” Load Model with Unsloth (QLoRA 4-bit)\n", + "> **Requires T4 GPU.** Skip to Cell 9 if you only want to test the environment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "model-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from unsloth import FastLanguageModel\n", + "\n", + "MODEL_NAME = \"unsloth/Qwen2.5-3B-Instruct\" # swap to 7B if VRAM allows\n", + "MAX_SEQ_LEN = 1024\n", + "LORA_RANK = 16\n", + "\n", + "model, tokenizer = FastLanguageModel.from_pretrained(\n", + " model_name = MODEL_NAME,\n", + " max_seq_length = MAX_SEQ_LEN,\n", + " dtype = None, # auto bf16/fp16\n", + " load_in_4bit = True,\n", + ")\n", + "\n", + "model = FastLanguageModel.get_peft_model(\n", + " model,\n", + " r = LORA_RANK,\n", + " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n", + " \"gate_proj\", \"up_proj\", \"down_proj\"],\n", + " lora_alpha = LORA_RANK * 2,\n", + " lora_dropout = 0.0,\n", + " bias = \"none\",\n", + " use_gradient_checkpointing = \"unsloth\",\n", + " random_state = 42,\n", + ")\n", + "\n", + "print(f\"βœ… Loaded {MODEL_NAME} with QLoRA rank={LORA_RANK}\")\n", + "print(f\" GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")" + ] + }, + { + "cell_type": "markdown", + "id": "prompt-header", + "metadata": {}, + "source": [ + "## πŸ’¬ Cell 6 β€” System Prompt & Dataset Builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "prompt-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "from datasets import Dataset\n", + "\n", + "\n", + "class Deliverable2_Formatter:\n", + " 'State translation + robust action-array extraction for Role 2.'\n", + "\n", + " @staticmethod\n", + " def translate_state(raw_json: list) -> str:\n", + " pseudo_assembly = []\n", + " for i, instruction in enumerate(raw_json):\n", + " if not isinstance(instruction, dict):\n", + " pseudo_assembly.append(f\"{i}. NOP\")\n", + " continue\n", + " op = str(instruction.get(\"op\", \"UNKNOWN\")).upper()\n", + " args = \", \".join(str(arg) for arg in instruction.get(\"args\", []))\n", + " dest = instruction.get(\"dest\", \"\")\n", + " if dest:\n", + " pseudo_assembly.append(f\"{i}. {dest} = {op} {args}\".rstrip())\n", + " else:\n", + " pseudo_assembly.append(f\"{i}. {op} {args}\".rstrip())\n", + " return \"\\n\".join(pseudo_assembly) if pseudo_assembly else \"; (empty program)\"\n", + "\n", + " @staticmethod\n", + " def extract_action_array(llm_output: str) -> list:\n", + " text = (llm_output or \"\").strip()\n", + " if not text:\n", + " raise ValueError(\"Invalid JSON format\")\n", + "\n", + " try:\n", + " parsed = json.loads(text)\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " pass\n", + "\n", + " match = re.search(r\"\\[.*?\\]\", text, re.DOTALL)\n", + " if match:\n", + " try:\n", + " parsed = json.loads(match.group(0))\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " pass\n", + "\n", + " raise ValueError(\"Invalid JSON format\")\n", + "\n", + "\n", + "def build_system_prompt(passes: dict) -> str:\n", + " lines: list = []\n", + " for name in sorted(passes.keys()):\n", + " if name == \"expand_constant\":\n", + " lines.append(\n", + " \" - expand_constant: Splits a constant assignment into runtime arithmetic (e.g., CONST 8 becomes CONST 3 + ADD 5). May temporarily increase cycles but can enable forward passes to find better optimization chains. Use sparingly when standard passes seem stuck.\"\n", + " )\n", + " elif name == \"duplicate_computation\":\n", + " lines.append(\n", + " \" - duplicate_computation: Creates a redundant copy of a binary operation with a fresh variable. May temporarily increase cycles but can enable alternative dead code elimination paths. Use when redundancy might unlock further simplification.\"\n", + " )\n", + " else:\n", + " lines.append(f\" - {name}\")\n", + " action_block = \"\\n\".join(lines)\n", + " return (\n", + " \"You are a compiler optimization agent. Your goal is to reduce \"\n", + " \"CPU cycle count by applying optimization passes to the program below.\\n\\n\"\n", + " f\"Available actions:\\n{action_block}\\n\"\n", + " \" - STOP (emit this in the array to stop; optional alias: \\\"done\\\". Stop early if no further improvement is possible.)\\n\\n\"\n", + " \"Rules:\\n\"\n", + " \" β€’ Output only a JSON array of pass names (example: [\\\"constant_folding\\\"]).\\n\"\n", + " \" β€’ No explanation, no markdown, no extra text.\\n\"\n", + " \" β€’ Do not invent actions not listed above.\\n\"\n", + " \" β€’ Maximum 5 passes per response (hard episode cap: 5 optimization steps or STOP).\\n\"\n", + " )\n", + "\n", + "\n", + "def build_dataset(\n", + " programs: list,\n", + " engine,\n", + " passes: dict,\n", + ") -> Dataset:\n", + " # Each row = one episode; GRPO samples K completions per row\n", + " env = CompilerOptimizationEnv(\n", + " engine, passes, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", + " )\n", + " system_prompt = build_system_prompt(passes)\n", + "\n", + " rows = []\n", + " for prog in programs:\n", + " translated_state = Deliverable2_Formatter.translate_state(prog)\n", + " prompt = [\n", + " {\"role\": \"system\", \"content\": system_prompt},\n", + " {\"role\": \"user\", \"content\": f\"Current program:\\n{translated_state}\\n\\nChoose optimization passes:\"},\n", + " ]\n", + " rows.append({\"prompt\": prompt, \"program_json\": json.dumps(prog)})\n", + "\n", + " return Dataset.from_list(rows)\n", + "\n", + "\n", + "# --- Demo: build dataset from mock programs -----------------------------------\n", + "def make_mock_programs(n: int = 20) -> list:\n", + " # Generate N random mock IR programs\n", + " import random\n", + " ops = [\"add\", \"mul\", \"sub\", \"const\", \"load\"]\n", + " progs = []\n", + " for _ in range(n):\n", + " length = random.randint(4, 12)\n", + " prog = [\n", + " {\n", + " \"op\": random.choice(ops),\n", + " \"dest\": f\"v{i}\",\n", + " \"args\": [f\"v{max(0, i-1)}\"],\n", + " \"type\": \"int\",\n", + " }\n", + " for i in range(length)\n", + " ]\n", + " prog.append({\"op\": \"ret\", \"args\": [f\"v{length-1}\"]})\n", + " progs.append(prog)\n", + " return progs\n", + "\n", + "\n", + "mock_programs = make_mock_programs(n=30)\n", + "train_dataset = build_dataset(mock_programs, MockEngine(), MOCK_PASSES)\n", + "\n", + "print(f\"βœ… Dataset built: {len(train_dataset)} episodes\")\n", + "print(f\" Sample prompt keys: {list(train_dataset[0].keys())}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "reward-fn-header", + "metadata": {}, + "source": [ + "## 🎯 Cell 7 β€” Reward Function Factory (for GRPO)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "reward-fn-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from typing import Any\n", + "\n", + "# Rollout + terminal weighting; reverse-pass W&B window logging\n", + "from toyir_rl_support import (\n", + " CompilerOptimizationEnv,\n", + " MOCK_PASSES,\n", + " rollout_shaped_return,\n", + " log_reverse_pass_stats_for_completion,\n", + " SAMPLE_PROGRAM,\n", + ")\n", + "\n", + "EP_CAP = CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", + "\n", + "\n", + "def _normalize_action_seq(actions: list, cap: int, passes: dict) -> list:\n", + " out: list = []\n", + " for x in actions[:cap]:\n", + " raw = str(x).strip()\n", + " s_low = raw.lower()\n", + " if s_low in (\"stop\", \"done\"):\n", + " out.append(\"done\" if s_low == \"done\" else \"STOP\")\n", + " continue\n", + " key = None\n", + " for k in passes:\n", + " if k.lower() == s_low:\n", + " key = k\n", + " break\n", + " if key is None:\n", + " out.append(raw)\n", + " else:\n", + " out.append(key)\n", + " return out\n", + "\n", + "\n", + "def make_reward_fn(engine, passes, max_steps: int = EP_CAP):\n", + " # TRL GRPO reward: terminal-weighted `rollout_shaped_return` + reverse-pass logging\n", + " _cap = min(int(max_steps), EP_CAP)\n", + "\n", + " def reward_fn(prompts, completions, **kwargs):\n", + " programs = kwargs.get(\"program_json\", [None] * len(completions))\n", + " rewards: list = []\n", + "\n", + " for completion, prog_json in zip(completions, programs):\n", + " raw = completion if isinstance(completion, str) else completion[0][\"content\"]\n", + " if prog_json is None:\n", + " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", + " continue\n", + " program = (\n", + " json.loads(prog_json) if isinstance(prog_json, str) else list(prog_json)\n", + " )\n", + " if not program:\n", + " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", + " continue\n", + " try:\n", + " actions = Deliverable2_Formatter.extract_action_array(raw)\n", + " except ValueError:\n", + " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", + " continue\n", + " if not actions:\n", + " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", + " continue\n", + " norm = _normalize_action_seq(actions, _cap, passes)\n", + " tr, n_e, n_d, _hs, _hc, ok = rollout_shaped_return(\n", + " program, norm, engine, passes\n", + " )\n", + " if not ok:\n", + " rewards.append(-1000.0)\n", + " else:\n", + " log_reverse_pass_stats_for_completion(n_e, n_d)\n", + " rewards.append(tr)\n", + " return rewards\n", + "\n", + " return reward_fn\n", + "\n", + "\n", + "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=EP_CAP)\n", + "print(\"βœ… Reward function factory ready\")\n", + "\n", + "# Quick sanity check\n", + "test_completions = [\n", + " '[\"constant_folding\"]',\n", + " 'Here is my plan: [\"dead_code_elimination\", \"peephole_optimization\", \"done\"]',\n", + " \"hallucinated_pass\",\n", + "]\n", + "test_programs = [json.dumps(SAMPLE_PROGRAM)] * 3\n", + "test_rewards = reward_fn(\n", + " prompts=[\"\"] * 3,\n", + " completions=test_completions,\n", + " program_json=test_programs,\n", + ")\n", + "print(\"\\nReward sanity check:\")\n", + "for act, rew in zip(test_completions, test_rewards):\n", + " print(f\" {act!r} β†’ {rew:+.2f}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "trainer-header", + "metadata": {}, + "source": [ + "## πŸš€ Cell 8 β€” GRPO Trainer Config & Training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "trainer-cell", + "metadata": {}, + "outputs": [], + "source": [ + "from trl import GRPOConfig, GRPOTrainer\n", + "\n", + "# Optional W&B β€” comment out if not using\n", + "try:\n", + " import wandb\n", + " wandb.init(project=\"openenv-compiler-opt\", name=\"grpo-qwen2.5-3b\")\n", + " REPORT_TO = \"wandb\"\n", + "except Exception:\n", + " REPORT_TO = \"none\"\n", + "\n", + "\n", + "grpo_config = GRPOConfig(\n", + " # ── Generation ────────────────────────────────────────────────────────\n", + " num_generations = 4, # K rollouts per prompt for group-relative advantage\n", + " max_new_tokens = 16, # Actions are 1 word; don't waste context\n", + " temperature = 0.9,\n", + " top_p = 0.95,\n", + "\n", + " # ── Optimisation ──────────────────────────────────────────────────────\n", + " learning_rate = 5e-6,\n", + " per_device_train_batch_size = 2,\n", + " gradient_accumulation_steps = 4, # effective batch = 8\n", + " num_train_epochs = 3,\n", + " max_grad_norm = 0.5,\n", + "\n", + " # ── GRPO-specific ─────────────────────────────────────────────────────\n", + " beta = 0.04, # KL penalty; keeps policy near reference\n", + "\n", + " # ── Logging / checkpointing ───────────────────────────────────────────\n", + " output_dir = \"./grpo_compiler_checkpoints\",\n", + " logging_steps = 10,\n", + " save_steps = 100,\n", + " report_to = REPORT_TO,\n", + "\n", + " # ── Reproducibility ───────────────────────────────────────────────────\n", + " seed = 42,\n", + ")\n", + "\n", + "trainer = GRPOTrainer(\n", + " model = model,\n", + " tokenizer = tokenizer,\n", + " config = grpo_config,\n", + " train_dataset = train_dataset,\n", + " reward_funcs = reward_fn,\n", + ")\n", + "\n", + "print(\"βœ… Trainer configured\")\n", + "print(f\" num_generations (K) = {grpo_config.num_generations}\")\n", + "print(f\" effective batch size = \"\n", + " f\"{grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n", + "print(f\" KL beta = {grpo_config.beta}\")\n", + "print()\n", + "print(\"Starting training... (this will take a while on T4)\")\n", + "trainer.train()" + ] + }, + { + "cell_type": "markdown", + "id": "save-header", + "metadata": {}, + "source": [ + "## πŸ’Ύ Cell 9 β€” Save Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "save-cell", + "metadata": {}, + "outputs": [], + "source": [ + "SAVE_PATH = \"./grpo_compiler_final\"\n", + "\n", + "model.save_pretrained(SAVE_PATH)\n", + "tokenizer.save_pretrained(SAVE_PATH)\n", + "\n", + "print(f\"βœ… Model saved to {SAVE_PATH}\")\n", + "\n", + "# Optional: push to HuggingFace Hub\n", + "# model.push_to_hub(\"your-hf-username/compiler-opt-grpo\")\n", + "# tokenizer.push_to_hub(\"your-hf-username/compiler-opt-grpo\")" + ] + }, + { + "cell_type": "markdown", + "id": "curriculum-header", + "metadata": {}, + "source": [ + "## πŸ“ˆ Cell 10 β€” Curriculum Callback & Reward Tracking" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "curriculum-cell", + "metadata": {}, + "outputs": [], + "source": [ + "class CurriculumCallback:\n", + " \"\"\"\n", + " Tracks rolling mean reward and promotes curriculum level\n", + " when the agent has mastered the current difficulty.\n", + "\n", + " Usage: call .record(reward) after every episode.\n", + " Read .level to get current difficulty (1=easy, 2=medium, 3=hard).\n", + " \"\"\"\n", + " def __init__(self, reward_threshold: float = 5.0, window: int = 50):\n", + " self.threshold = reward_threshold\n", + " self.window = window\n", + " self._history = []\n", + " self.level = 1\n", + " self._promotions = []\n", + "\n", + " def record(self, reward: float, step: int = None):\n", + " self._history.append(reward)\n", + " if len(self._history) >= self.window:\n", + " mean = sum(self._history[-self.window:]) / self.window\n", + " if mean >= self.threshold and self.level < 3:\n", + " self.level += 1\n", + " self._promotions.append((step or len(self._history), self.level))\n", + " print(f\"[Curriculum] β–² Promoted to level {self.level} \"\n", + " f\"(rolling mean={mean:.2f})\")\n", + "\n", + " def plot(self):\n", + " import matplotlib.pyplot as plt\n", + " import numpy as np\n", + "\n", + " history = self._history\n", + " steps = list(range(len(history)))\n", + " window = self.window\n", + " rolling = [\n", + " sum(history[max(0,i-window):i+1]) / min(i+1, window)\n", + " for i in steps\n", + " ]\n", + "\n", + " fig, ax = plt.subplots(figsize=(10, 4))\n", + " ax.plot(steps, history, alpha=0.3, color=\"#3498db\", label=\"Episode reward\")\n", + " ax.plot(steps, rolling, color=\"#e74c3c\", linewidth=2,\n", + " label=f\"Rolling mean (w={window})\")\n", + " ax.axhline(self.threshold, linestyle=\"--\", color=\"grey\",\n", + " linewidth=1, label=f\"Promotion threshold ({self.threshold})\")\n", + " for step, level in self._promotions:\n", + " ax.axvline(step, color=\"green\", linewidth=1.5, linestyle=\":\")\n", + " ax.text(step, ax.get_ylim()[1]*0.9, f\" L{level}\",\n", + " color=\"green\", fontsize=9)\n", + " ax.set_xlabel(\"Episode\")\n", + " ax.set_ylabel(\"Reward\")\n", + " ax.set_title(\"Training Reward + Curriculum Progression\")\n", + " ax.legend()\n", + " plt.tight_layout()\n", + " plt.savefig(\"curriculum_reward_curve.png\", dpi=150)\n", + " plt.show()\n", + " print(\"πŸ“ˆ Saved curriculum_reward_curve.png\")\n", + "\n", + "\n", + "# ── Demo: simulate 200 episodes of improving reward ──────────────────────────\n", + "import random\n", + "cb = CurriculumCallback(reward_threshold=5.0, window=50)\n", + "for ep in range(200):\n", + " # Simulate reward slowly improving\n", + " synthetic_reward = -5 + ep * 0.08 + random.gauss(0, 2)\n", + " cb.record(synthetic_reward, step=ep)\n", + "\n", + "cb.plot()" + ] + }, + { + "cell_type": "markdown", + "id": "inference-header", + "metadata": {}, + "source": [ + "## πŸ” Cell 11 β€” Inference: Before vs After Training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "inference-cell", + "metadata": {}, + "outputs": [], + "source": [ + "def run_inference_episode(model, tokenizer, engine, passes, program, max_steps=5):\n", + " \"\"\"Run inference and execute parsed pass arrays until episode termination.\"\"\"\n", + " FastLanguageModel.for_inference(model)\n", + "\n", + " _ms = min(int(max_steps), CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n", + " env = CompilerOptimizationEnv(engine, passes, max_steps=_ms)\n", + " obs = env.reset(program)\n", + " system_prompt = build_system_prompt(passes)\n", + "\n", + " actions_chosen, rewards_earned = [], []\n", + " done = False\n", + "\n", + " print(f\"\\nBaseline cycles: {env.previous_cycles}\")\n", + " print(f\"Initial state:\\n{obs}\\n\")\n", + "\n", + " while not done:\n", + " messages = [\n", + " {\"role\": \"system\", \"content\": system_prompt},\n", + " {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose optimization passes:\"},\n", + " ]\n", + " inputs = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=True,\n", + " add_generation_prompt=True,\n", + " return_tensors=\"pt\",\n", + " ).to(model.device)\n", + "\n", + " with torch.no_grad():\n", + " outputs = model.generate(\n", + " input_ids=inputs,\n", + " max_new_tokens=32,\n", + " temperature=0.1,\n", + " do_sample=True,\n", + " )\n", + "\n", + " raw_output = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip()\n", + "\n", + " try:\n", + " parsed_actions = Deliverable2_Formatter.extract_action_array(raw_output)\n", + " except ValueError:\n", + " parsed_actions = []\n", + "\n", + " if not parsed_actions:\n", + " result = env.step(\"__invalid__\")\n", + " actions_chosen.append(\"__invalid__\")\n", + " rewards_earned.append(result.reward)\n", + " obs = result.observation\n", + " done = result.done\n", + " print(f\"Step {len(actions_chosen)}: invalid output '{raw_output}' β†’ reward={result.reward:+.2f}\")\n", + " continue\n", + "\n", + " for action in parsed_actions[:CompilerOptimizationEnv.MAX_EPISODE_STEPS]:\n", + " a = str(action).strip()\n", + " al = a.lower()\n", + " if al in (\"done\", \"stop\"):\n", + " result = env.step(\"done\" if al == \"done\" else \"STOP\")\n", + " actions_chosen.append(al)\n", + " rewards_earned.append(result.reward)\n", + " obs = result.observation\n", + " done = result.done\n", + " print(f\"Step {len(actions_chosen)}: stop ('{a}') β†’ reward={result.reward:+.2f}\")\n", + " break\n", + " action = al\n", + " if action not in passes:\n", + " for k in passes:\n", + " if k.lower() == al:\n", + " action = k\n", + " break\n", + " result = env.step(action)\n", + " actions_chosen.append(action)\n", + " rewards_earned.append(result.reward)\n", + " obs = result.observation\n", + " done = result.done\n", + " print(f\"Step {len(actions_chosen)}: '{action}' β†’ reward={result.reward:+.2f}\")\n", + " if done:\n", + " break\n", + "\n", + " summary = env._episode_summary()\n", + " print(f\"\\n{'─'*40}\")\n", + " print(f\"Total improvement: {summary['total_improvement_pct']:.1f}%\")\n", + " print(f\"Final cycles: {summary['final_cycles']} (was {summary['baseline_cycles']})\")\n", + " return summary\n", + "\n", + "\n", + "# Uncomment after training:\n", + "# summary = run_inference_episode(\n", + "# model, tokenizer, MockEngine(), MOCK_PASSES, SAMPLE_PROGRAM\n", + "# )\n", + "\n", + "print(\"βœ… Inference cell ready. Uncomment the last block after training to run.\")" + ] + }, + { + "cell_type": "markdown", + "id": "tips-header", + "metadata": {}, + "source": [ + "---\n", + "## πŸ“ Notes & Tips\n", + "\n", + "| What | Why it matters |\n", + "|------|----------------|\n", + "| `num_generations=4` | GRPO needs Kβ‰₯2 rollouts per prompt to compute group-relative advantage. K=4 balances diversity vs. compute. |\n", + "| `beta=0.04` | KL penalty keeping policy close to reference. Too low β†’ mode collapse. Too high β†’ no learning. |\n", + "| `max_new_tokens=16` | Actions are one word. This prevents wasted computation and keeps the model from adding explanations. |\n", + "| Equivalence / semantic failure | The shaped reward function returns `βˆ’1000.0` when a candidate fails verification (RLVR hard gate). |\n", + "| Terminal reward | Final cycle savings (scaled) are given at `STOP`, step cap (5), or at end of a pass list; non-terminal steps use a small constant cost. |\n", + "| Soft invalid-action termination | 3 consecutive invalid actions β†’ end. Single mistakes don't kill the episode; the agent can recover. |\n", + "| `deepcopy` on all pass inputs | Role 3's passes mutate dicts in-place. Without this, `original_program` gets corrupted and the verifier fails spuriously. |\n", + "\n", + "**Next steps:**\n", + "- Replace `MockEngine` with Role 1's real engine\n", + "- Replace `MOCK_PASSES` with Role 3's real passes\n", + "- Push environment to HuggingFace Spaces: `openenv init && openenv deploy`\n", + "- Add W&B sweep to tune `beta`, `learning_rate`, `num_generations`" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/space/space/space/space/space/space/space/reverse_pass/toyir_rl_support.py b/space/space/space/space/space/space/space/reverse_pass/toyir_rl_support.py new file mode 100644 index 0000000000000000000000000000000000000000..c347da8e95a36b1e3d36e89cc39728454faebf81 --- /dev/null +++ b/space/space/space/space/space/space/space/reverse_pass/toyir_rl_support.py @@ -0,0 +1,473 @@ +""" +Toy-IR RL training support: PASSES, mock engine, environment, terminal-weighted shaped reward. +Used by training notebooks. Do not import toy_vm / verifier / reward_utils here. +""" + +from __future__ import annotations + +import copy +import json +import os +import sys +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +# Allow notebooks in subdirs to import when cwd is set to project root +_ROOT = os.path.dirname(os.path.abspath(__file__)) + + +@dataclass +class StepResult: + observation: str + reward: float + done: bool + info: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class EpisodeStats: + steps_taken: int = 0 + total_reward: float = 0.0 + passes_applied: List[str] = field(default_factory=list) + invalid_actions: int = 0 + no_ops: int = 0 + baseline_cycles: int = 0 + final_cycles: int = 0 + + @property + def total_improvement_pct(self) -> float: + if self.baseline_cycles == 0: + return 0.0 + return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0 + + +class MCPEnvironment: + def reset(self, *args, **kwargs): + raise NotImplementedError + + def step(self, *args, **kwargs): + raise NotImplementedError + + def state(self): + raise NotImplementedError + + +def _op_u(ins: dict) -> str: + return str(ins.get("op", "")).upper() + + +def _first_numeric(ins: dict) -> Optional[int]: + for k in ("src1", "src2"): + v = ins.get(k) + if isinstance(v, int) and v > 0: + return v + args = ins.get("args", []) + if args: + a0 = args[0] + try: + if isinstance(a0, int): + return a0 + return int(str(a0), 0) + except (ValueError, TypeError): + return None + v = ins.get("src1") + if isinstance(v, int): + return v + return None + + +# --- forward / reverse passes (list[dict] IR) --------------------------------- + + +def constant_folding(program: List[dict]) -> List[dict]: + p = copy.deepcopy(program) + return p[:-1] if len(p) > 1 else p + + +def dead_code_elimination(program: List[dict]) -> List[dict]: + p = copy.deepcopy(program) + return p[:-1] if len(p) > 2 else p + + +def peephole_optimization(program: List[dict]) -> List[dict]: + return copy.deepcopy(program) + + +def expand_constant(program: List[dict]) -> List[dict]: + """ + Deoptimization: split one CONST into two consts + ADD to raise cycle count sometimes, + giving forward passes a different surface (reverse exploration). + """ + p = copy.deepcopy(program) + for i, ins in enumerate(p): + if _op_u(ins) not in ("CONST",): + continue + v = _first_numeric(ins) if _first_numeric(ins) is not None else None + if v is None and ins.get("args"): + try: + v = int(ins["args"][0]) + except (ValueError, TypeError, IndexError, KeyError): + v = None + if v is None or v < 2: + continue + d = str(ins.get("dest", "t0")) + h, rest = v // 2, v - (v // 2) + repl = [ + {**{k: ins[k] for k in ins if k in ("type",)}, "op": "const", "dest": f"{d}_a", "args": [str(h)]}, + {**{k: ins[k] for k in ins if k in ("type",)}, "op": "const", "dest": f"{d}_b", "args": [str(rest)]}, + { + "op": "add", + "dest": d, + "args": [f"{d}_a", f"{d}_b"], + **{k: ins[k] for k in ins if k in ("type",) and k not in ("op", "dest", "args")}, + }, + ] + p[i : i + 1] = repl + return p + return p + + +def duplicate_computation(program: List[dict]) -> List[dict]: + """ + Deoptimization: duplicate a binary op into a new destination (redundant work). + """ + p = copy.deepcopy(program) + for i, ins in enumerate(p): + ou = _op_u(ins) + if ou in ("ADD", "MUL", "SUB", "DIV"): + dest = str(ins.get("dest", "t0")) + dup = copy.deepcopy(ins) + dup["dest"] = f"{dest}_d" + p.insert(i + 1, dup) + return p + return p + + +PASSES: Dict[str, Any] = { + "constant_folding": constant_folding, + "dead_code_elimination": dead_code_elimination, + "peephole_optimization": peephole_optimization, + "expand_constant": expand_constant, + "duplicate_computation": duplicate_computation, +} + +MOCK_PASSES = PASSES # alias for existing notebook references + + +class MockEngine: + """Stub: cycle count = instruction list length; always equivalent.""" + + def execute_and_count_cycles(self, program: List[dict]) -> int: + return len(program) + + def verify_equivalence(self, original, candidate) -> bool: + return True + + +# Back-compat: some notebooks use these names +MOCK_ENGINE = MockEngine() + + +# --- Shaped reward (user-specified) ------------------------------------------- + + +def compute_shaped_reward( + equivalent: bool, + is_terminal: bool, + original_cycles: int, + current_cycles: int, +) -> float: + # RLVR hard gate β€” broken optimizations get massive penalty + if not equivalent: + return -1000.0 + + # Terminal reward: large signal proportional to final cycle reduction + # This is what makes reverse passes learnable β€” agent gets rewarded for the END state, not the path + if is_terminal: + if original_cycles == 0: + return 0.0 + relative_savings = (original_cycles - current_cycles) / original_cycles + return relative_savings * 100.0 + + # Per-step: tiny constant cost to discourage infinite thrashing + # CRITICAL: do NOT punish cycle increases here. Per-step cycle-delta reward kills reverse-pass exploration. + return -0.1 + + +# --- Environment -------------------------------------------------------------- + + +class CompilerOptimizationEnv(MCPEnvironment): + TIME_TAX: float = 1.0 + NO_OP_PENALTY: float = -2.0 + INVALID_ACTION_PENALTY: float = -5.0 + MAX_INVALID_ACTIONS: int = 3 + TERMINAL_BONUS_SCALE: float = 0.5 + MAX_EPISODE_STEPS: int = 5 # hard cap (Change 4) + + def __init__( + self, + role1_engine: Any, + role3_passes: Dict[str, Any], + max_steps: int = 5, + curriculum_level: int = 1, + ): + self.engine = role1_engine + self.passes = role3_passes + self.max_steps = max(1, int(max_steps)) + self.curriculum_level = curriculum_level + self._valid_actions: frozenset = frozenset(self.passes.keys()) | {"STOP", "done", "DONE"} + self._stats: Optional[EpisodeStats] = None + self.original_program: Optional[List[dict]] = None + self.current_program: Optional[List[dict]] = None + self.previous_cycles = 0 + self._consecutive_invalid = 0 + + def reset(self, new_program_json: List[dict]) -> str: + self.original_program = copy.deepcopy(new_program_json) + self.current_program = copy.deepcopy(new_program_json) + self.previous_cycles = self._safe_count_cycles(self.current_program) + self._consecutive_invalid = 0 + self._stats = EpisodeStats( + baseline_cycles=self.previous_cycles, + final_cycles=self.previous_cycles, + ) + return self.state() + + def state(self) -> str: + return self._program_to_pseudoasm(self.current_program) + + def step(self, action_string: str) -> StepResult: + assert self._stats is not None, "Call reset() before step()." + aup = str(action_string).upper() + is_stop = aup in ("STOP", "DONE") + if is_stop: + self._stats.steps_taken += 1 + orig = int(self._stats.baseline_cycles) + cur = self._safe_count_cycles(self.current_program) + ok = bool(self.engine.verify_equivalence(self.original_program, self.current_program)) + reward = compute_shaped_reward(ok, True, orig, cur) + self._stats.total_reward += reward + return StepResult( + self.state(), + reward, + True, + { + "reason": "stop", + "is_terminal": True, + "stop_token": aup, + "episode_stats": self._episode_summary(), + }, + ) + if action_string not in self.passes: + return self._handle_invalid_action(action_string) + + self._stats.steps_taken += 1 + candidate = self.passes[action_string](copy.deepcopy(self.current_program)) + if not self.engine.verify_equivalence(self.original_program, candidate): + return self._handle_semantic_violation() + new_cycles = self._safe_count_cycles(candidate) + reward, info = self._compute_reward(action_string, new_cycles) + self.current_program = candidate + self.previous_cycles = new_cycles + self._stats.final_cycles = new_cycles + self._stats.total_reward += reward + self._stats.passes_applied.append(action_string) + self._consecutive_invalid = 0 + done = self._stats.steps_taken >= self.max_steps + if done: + info["reason"] = "max_steps_reached" + info["episode_stats"] = self._episode_summary() + return StepResult(self.state(), reward, done, info) + + def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, dict]: + info: Dict[str, Any] = {"action": action} + orig = int(self._stats.baseline_cycles) + is_term = self._stats.steps_taken >= self.max_steps + if new_cycles == self.previous_cycles and not is_term: + self._stats.no_ops += 1 + info["no_op"] = True + reward = compute_shaped_reward(True, is_term, orig, new_cycles) + info["is_terminal"] = is_term + info["new_cycles"] = new_cycles + return reward, info + + def _compute_crash_penalty(self) -> float: + return -1000.0 + + def _handle_invalid_action(self, action: str) -> StepResult: + self._consecutive_invalid += 1 + self._stats.invalid_actions += 1 + done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS + info: Dict[str, Any] = { + "error": f"Unknown action: '{action}'", + "valid_actions": sorted(self._valid_actions), + "consecutive_invalid": self._consecutive_invalid, + } + if done: + info["reason"] = "too_many_invalid_actions" + info["episode_stats"] = self._episode_summary() + return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info) + + def _handle_semantic_violation(self) -> StepResult: + return StepResult( + self.state(), + self._compute_crash_penalty(), + True, + { + "error": "Semantic equivalence check FAILED.", + "reason": "semantic_violation", + "episode_stats": self._episode_summary(), + }, + ) + + @staticmethod + def _program_to_pseudoasm(program: Optional[List[dict]]) -> str: + if not program: + return "; (empty program)" + lines = [] + for i, instr in enumerate(program): + op = instr.get("op", "NOP") + args = instr.get("args", []) + dest = instr.get("dest") + typ = instr.get("type", "") + arg_str = ", ".join(str(a) for a in args) + type_hint = f":{typ}" if typ else "" + if dest: + lines.append(f" {i:>3}: {dest}{type_hint} = {op} {arg_str}") + else: + lines.append(f" {i:>3}: {op} {arg_str}") + return "\n".join(lines) + + def _safe_count_cycles(self, program: List[dict]) -> int: + return max(0, int(self.engine.execute_and_count_cycles(program))) + + def _episode_summary(self) -> dict: + s = self._stats + return { + "steps": s.steps_taken, + "total_reward": round(s.total_reward, 3), + "passes_applied": s.passes_applied, + "invalid_actions": s.invalid_actions, + "no_ops": s.no_ops, + "baseline_cycles": s.baseline_cycles, + "final_cycles": s.final_cycles, + "total_improvement_pct": round(s.total_improvement_pct, 3), + } + + def available_actions(self) -> List[str]: + return sorted(self._valid_actions) + + +SAMPLE_PROGRAM: List[dict] = [ + {"op": "const", "dest": "x", "args": ["5"], "type": "int"}, + {"op": "const", "dest": "y", "args": ["3"], "type": "int"}, + {"op": "add", "dest": "z", "args": ["x", "y"], "type": "int"}, + {"op": "mul", "dest": "w", "args": ["z", "x"], "type": "int"}, + {"op": "ret", "args": ["w"]}, +] + + +# --- Rollout for GRPO reward (terminal-weighted shaping) ------------------------ + + +def rollout_shaped_return( + program: List[dict], + action_names: List[str], + engine: Any, + passes: Optional[Dict[str, Any]] = None, +) -> Tuple[float, int, int, bool, bool, bool]: + """ + Returns: + (total_shaped_reward, n_expand_constant, n_duplicate, hit_stop, hit_step_cap, parse_ok) + If parse_ok is False (unknown pass in sequence), first value is -1000.0. + Sums per-step `env.step` shaped rewards, adds a terminal `compute_shaped_reward` when + the action list ends without an explicit STOP (or already terminal from step cap). + """ + passes = passes or PASSES + cap = CompilerOptimizationEnv.MAX_EPISODE_STEPS + env = CompilerOptimizationEnv(engine, passes, max_steps=cap) + env.reset(program) + total = 0.0 + n_e = 0 + n_d = 0 + hit_stop = False + hit_cap = False + for raw in action_names: + a = str(raw).strip() + if not a: + continue + aup = a.upper() + if aup in ("STOP", "DONE"): + r = env.step(a) + total += r.reward + hit_stop = True + return total, n_e, n_d, hit_stop, hit_cap, True + if a not in passes: + return -1000.0, n_e, n_d, hit_stop, hit_cap, False + if a == "expand_constant": + n_e += 1 + elif a == "duplicate_computation": + n_d += 1 + r = env.step(a) + total += r.reward + if r.info.get("reason") == "semantic_violation": + return -1000.0, n_e, n_d, hit_stop, hit_cap, True + if r.done and r.info.get("reason") == "max_steps_reached": + hit_cap = True + return total, n_e, n_d, hit_stop, hit_cap, True + ok = bool(engine.verify_equivalence(env.original_program, env.current_program)) + cur = int(engine.execute_and_count_cycles(env.current_program or [])) + orig = int(env._stats.baseline_cycles) + total += compute_shaped_reward(ok, True, orig, cur) + return total, n_e, n_d, hit_stop, hit_cap, True + + +# Global window for logging (50-episode / completion windows) +_REVERSE_LOG_WINDOW: List[dict] = [] + + +def log_reverse_pass_stats_for_completion( + used_expand: int, + used_dup: int, +) -> None: + """Call once per training completion. Logs every 50 'episodes' (completions).""" + global _REVERSE_LOG_WINDOW + any_r = (used_expand + used_dup) > 0 + _REVERSE_LOG_WINDOW.append( + { + "any_reverse": any_r, + "expand": used_expand, + "dup": used_dup, + } + ) + if len(_REVERSE_LOG_WINDOW) < 50: + return + w = _REVERSE_LOG_WINDOW + _REVERSE_LOG_WINDOW = [] + reverse_pass_episodes = sum(1 for e in w if e["any_reverse"]) + expand_constant_count = sum(e["expand"] for e in w) + duplicate_computation_count = sum(e["dup"] for e in w) + payload = { + "reverse_pass_episodes": reverse_pass_episodes, + "expand_constant_count": expand_constant_count, + "duplicate_computation_count": duplicate_computation_count, + } + try: + import wandb + + if wandb.run is not None: + wandb.log(payload) + except Exception: + pass + print( + f"[reverse_pass/50] reverse_pass_episodes={reverse_pass_episodes} " + f"expand_constant_count={expand_constant_count} " + f"duplicate_computation_count={duplicate_computation_count}" + ) + + +def ensure_import_path() -> None: + d = _ROOT + if d and d not in sys.path: + sys.path.insert(0, d) diff --git a/space/space/space/space/space/space/space/space/reverse_pass/reversepass_new_eval_baseline.ipynb b/space/space/space/space/space/space/space/space/reverse_pass/reversepass_new_eval_baseline.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..54917e9157c1e83629b45375db8691f737adbdc0 --- /dev/null +++ b/space/space/space/space/space/space/space/space/reverse_pass/reversepass_new_eval_baseline.ipynb @@ -0,0 +1,3570 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "code", + "source": [ + "## Design Assumptions (do not violate)\n", + "\n", + "# 1. **DCE never eliminates STOREs.** They define program output (final mem state).\n", + "# 2. **Addresses are distinct by construction.** Generator allocates each address variable to a unique integer; no aliasing.\n", + "# 3. **CF refuses to fold DIV by zero.** If src2 == 0 on a DIV op, leave instruction unchanged.\n", + "# 4. **Generator never emits DIV by literal zero.** When DIV is generated, src2 is always a non-zero constant or a variable known to be non-zero.\n", + "# # 5. **Integer arithmetic only.** No floats anywhere β€” avoids equivalence-check precision issues.\n", + "# 6. Generator declares `observable_addrs` per program β€” verifier compares only these mem entries.\n", + "# 7. State translator annotates observable outputs at top of dump.\n", + "# 8. Reward distinguishes broken (-1000) from valid-but-worse (small negative) β€” Harshal's formula.\n", + "# 9. Multi-input verification: 3-5 random initial states, all must match.\n", + "# 10. Integer division uses Python floor division (//). Aarush's VM must match." + ], + "metadata": { + "id": "_XI5jT2Ibvrf" + }, + "execution_count": 3, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "id": "j6CQ327KWXwr" + }, + "outputs": [], + "source": [ + "# === TAC Schema v1.0 (LOCKED with Role 1 / Aarush) ===\n", + "# Reverse passes deferred to stretch goal β€” not in initial action space.\n", + "\n", + "OPS = [\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\", \"STORE\", \"NOP\"]\n", + "\n", + "CYCLE_COSTS = {\n", + " \"CONST\": 1,\n", + " \"ADD\": 1,\n", + " \"SUB\": 1,\n", + " \"MUL\": 3,\n", + " \"DIV\": 5,\n", + " \"LOAD\": 4,\n", + " \"STORE\": 4,\n", + " \"NOP\": 0,\n", + "}\n", + "\n", + "# Instruction shape: {\"op\": str, \"dest\": str|None, \"src1\": Any, \"src2\": Any}\n", + "# Operands: str = variable name, int = literal constant, None = unused\n", + "#\n", + "# Op semantics:\n", + "# CONST: dest = src1 (src1 is int literal, src2 = None)\n", + "# ADD/SUB/MUL/DIV: dest = src1 OP src2 (src1, src2 are var names or int literals)\n", + "# LOAD: dest = mem[src1] (src1 is a var holding an address)\n", + "# STORE: mem[dest] = src1 (dest is a var holding an address)\n", + "# NOP: no-op (all fields None)\n", + "#\n", + "# Program output (for equivalence check) = final memory state (mem dict).\n", + "# Programs ship as: (initial_vars: dict, initial_mem: dict, instructions: list[dict])" + ] + }, + { + "cell_type": "code", + "source": [ + "import random\n", + "\n", + "def generate_level_1():\n", + " \"\"\"\n", + " Generate a Level 1 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 4-6 instructions\n", + " - 2-3 CONST ops with literal values\n", + " - 1-2 arithmetic ops on those constants (foldable by CF)\n", + " - 0-1 dead variables (killable by DCE)\n", + " - Exactly 1 STORE at the end so the program has an observable output\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " # Step 1: Generate 2-3 constant assignments\n", + " num_consts = random.randint(2, 3)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: Generate 1-2 arithmetic ops using those constants\n", + " num_arith = random.randint(1, 2)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"MUL\"])\n", + " src1 = random.choice(const_vars)\n", + " src2 = random.choice(const_vars)\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " last_result = dest\n", + " const_vars.append(dest)\n", + "\n", + " # Step 3: Optionally add 1 dead variable (50% chance)\n", + " if random.random() < 0.5:\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None})\n", + " # Note: dead_var is intentionally never used β€” DCE should catch it.\n", + "\n", + " # Step 4: Add a STORE at the end so the program has observable output\n", + " initial_vars = {\"addr0\": 0}\n", + " instructions.append({\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None})\n", + "\n", + " initial_mem = {}\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": [0], # Aarush's verifier compares only these mem entries\n", + " }\n", + "\n", + "\n", + "# Sanity-check: generate a few programs and print them\n", + "for seed in [42, 1, 7, 99]:\n", + " random.seed(seed)\n", + " prog = generate_level_1()\n", + " print(f\"\\n=== seed={seed} ===\")\n", + " print(f\"initial_vars : {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + " print(f\"observable_addrs : {prog['observable_addrs']}\")\n", + " print(f\"instructions ({len(prog['instructions'])}):\")\n", + " for i, instr in enumerate(prog['instructions']):\n", + " print(f\" {i}: {instr}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "a86vD1OyWcB9", + "outputId": "e94a597c-d305-4b8f-90fa-5d26e3cd84bf" + }, + "execution_count": 5, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== seed=42 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (4):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "=== seed=1 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 2, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v1', 'src2': 'v1'}\n", + " 3: {'op': 'MUL', 'dest': 'v3', 'src1': 'v2', 'src2': 'v1'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=7 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (6):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=99 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 4, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 10, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v0', 'src2': 'v0'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def _resolve_operand(operand, known_constants):\n", + " \"\"\"\n", + " Given a TAC operand (string variable name or int literal),\n", + " return its concrete integer value if known, or None if unknown.\n", + " \"\"\"\n", + " if isinstance(operand, int):\n", + " return operand\n", + " if isinstance(operand, str) and operand in known_constants:\n", + " return known_constants[operand]\n", + " return None\n", + "\n", + "\n", + "def _compute(op, v1, v2):\n", + " \"\"\"Compute the result of a binary arithmetic op on two known integers.\"\"\"\n", + " if op == \"ADD\": return v1 + v2\n", + " if op == \"SUB\": return v1 - v2\n", + " if op == \"MUL\": return v1 * v2\n", + " if op == \"DIV\": return v1 // v2 # floor division (locked semantics)\n", + " raise ValueError(f\"_compute called with non-arithmetic op: {op}\")\n", + "\n", + "\n", + "def constant_folding(program):\n", + " \"\"\"\n", + " Forward pass that folds constant arithmetic into CONST ops, and\n", + " propagates known constants into instruction operands.\n", + "\n", + " Behavior:\n", + " - If both sources of an arithmetic op resolve to known integers,\n", + " replaces the instruction with a CONST holding the computed result.\n", + " - If only one source is known, still substitutes that known value\n", + " into the instruction (constant propagation), enabling downstream\n", + " passes (e.g., peephole) to recognize patterns like ADD-with-0 or MUL-by-1.\n", + " - Refuses to fold DIV by zero (Design Assumption #3).\n", + " - Always returns fresh dicts; never aliases input instructions.\n", + "\n", + " Args:\n", + " program: list of TAC instruction dicts (per locked schema).\n", + "\n", + " Returns:\n", + " new list of TAC instruction dicts. Always semantics-preserving.\n", + " Never raises, never returns None.\n", + " \"\"\"\n", + " known_constants = {}\n", + " new_program = []\n", + "\n", + " for instr in program:\n", + " # Always work on a copy β€” never alias input dicts\n", + " instr = instr.copy()\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + "\n", + " if op == \"CONST\":\n", + " known_constants[dest] = instr[\"src1\"]\n", + " new_program.append(instr)\n", + "\n", + " elif op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n", + " # === Constant propagation: substitute known constants into operands ===\n", + " if isinstance(instr[\"src1\"], str) and instr[\"src1\"] in known_constants:\n", + " instr[\"src1\"] = known_constants[instr[\"src1\"]]\n", + " if isinstance(instr[\"src2\"], str) and instr[\"src2\"] in known_constants:\n", + " instr[\"src2\"] = known_constants[instr[\"src2\"]]\n", + "\n", + " # === Try to fold ===\n", + " v1 = _resolve_operand(instr[\"src1\"], known_constants)\n", + " v2 = _resolve_operand(instr[\"src2\"], known_constants)\n", + "\n", + " if v1 is not None and v2 is not None:\n", + " # Both operands are known integers\n", + " if op == \"DIV\" and v2 == 0:\n", + " # Refuse to fold DIV by zero\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + " else:\n", + " # Fold: replace with CONST\n", + " result = _compute(op, v1, v2)\n", + " new_program.append({\n", + " \"op\": \"CONST\",\n", + " \"dest\": dest,\n", + " \"src1\": result,\n", + " \"src2\": None,\n", + " })\n", + " known_constants[dest] = result\n", + " else:\n", + " # Can't fold (at least one operand unknown).\n", + " # Instruction may still have been mutated by propagation above.\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op == \"LOAD\":\n", + " # Memory reads aren't statically resolvable\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op in (\"STORE\", \"NOP\"):\n", + " # No dest tracking needed.\n", + " # Note: we COULD propagate src1 into a STORE for downstream readability,\n", + " # but the cycle cost is unchanged and the executor handles vars fine.\n", + " # Leave STORE alone β€” keeps the code minimal.\n", + " new_program.append(instr)\n", + "\n", + " else:\n", + " # Unknown op β€” defensive pass-through\n", + " new_program.append(instr)\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "CGk8Fz4CZA3t" + }, + "execution_count": 6, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for constant_folding ===\n", + "\n", + "# Test 1: simple fold β€” ADD of two CONSTs\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result1 = constant_folding(test1)\n", + "print(\"Test 1 (simple ADD fold):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 becomes CONST c = 8\n", + "\n", + "# Test 2: chained fold β€” second op uses first op's folded result\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # β†’ 8\n", + " {\"op\": \"MUL\", \"dest\": \"d\", \"src1\": \"c\", \"src2\": \"b\"}, # β†’ 8 * 5 = 40\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"d\", \"src2\": None},\n", + "]\n", + "result2 = constant_folding(test2)\n", + "print(\"\\nTest 2 (chained fold):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: c becomes CONST 8, d becomes CONST 40\n", + "\n", + "# Test 3: DIV by zero refused\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 10, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"DIV\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # 10/0 β€” must NOT fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result3 = constant_folding(test3)\n", + "print(\"\\nTest 3 (DIV by zero refused):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still DIV, not CONST)\n", + "\n", + "# Test 4: unknown source can't fold\n", + "test4 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b is unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # can't fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result4 = constant_folding(test4)\n", + "print(\"\\nTest 4 (LOAD makes b unknown, ADD not folded):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still ADD)\n", + "\n", + "# Test 5: idempotence β€” running CF on a generated program\n", + "random.seed(42)\n", + "prog = generate_level_1()\n", + "folded = constant_folding(prog[\"instructions\"])\n", + "print(\"\\nTest 5 (CF on generated Level 1 program, seed=42):\")\n", + "print(\"Before:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "print(\"After:\")\n", + "for i, instr in enumerate(folded): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 6: constant propagation β€” only one source is known\n", + "test6 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # a known (=0), b unknown\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result6 = constant_folding(test6)\n", + "print(\"\\nTest 6 (propagation: a=0 substituted into ADD even though b unknown):\")\n", + "for i, instr in enumerate(result6): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 is still ADD (can't fold β€” b unknown), but src1 is now literal 0, not 'a'\n", + "# This sets up peephole to recognize \"ADD with 0\" later" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "WLKSTVa8fRF8", + "outputId": "99f8dd11-fcf5-4598-ab03-6688af799bc0" + }, + "execution_count": 7, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (simple ADD fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 2 (chained fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'd', 'src1': 40, 'src2': None}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'd', 'src2': None}\n", + "\n", + "Test 3 (DIV by zero refused):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'DIV', 'dest': 'c', 'src1': 10, 'src2': 0}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 4 (LOAD makes b unknown, ADD not folded):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 3, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 5 (CF on generated Level 1 program, seed=42):\n", + "Before:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "Test 6 (propagation: a=0 substituted into ADD even though b unknown):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 0, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 0, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def dead_code_elimination(program):\n", + " \"\"\"\n", + " Backward liveness analysis. Replace dead instructions with NOPs,\n", + " then strip NOPs.\n", + "\n", + " Rules:\n", + " - STOREs are always live (define program output).\n", + " - Address variables used in STORE/LOAD are always live.\n", + " - Any instruction whose dest is never read later is dead.\n", + "\n", + " Returns a fresh list of instruction dicts. Never raises, never returns None.\n", + " \"\"\"\n", + " # Walk backward, build new program in reverse, then re-reverse at the end\n", + " live = set()\n", + " new_program_reversed = []\n", + "\n", + " for instr in reversed(program):\n", + " instr = instr.copy() # never alias input\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + "\n", + " if op == \"STORE\":\n", + " # STOREs are always kept; mark sources live\n", + " if isinstance(src1, str): live.add(src1)\n", + " if isinstance(dest, str): live.add(dest) # address variable\n", + " new_program_reversed.append(instr)\n", + "\n", + " elif op == \"NOP\":\n", + " # Pass through; will be stripped at the end\n", + " new_program_reversed.append(instr)\n", + "\n", + " elif op in (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\"):\n", + " if dest in live:\n", + " # Live instruction β€” keep it, mark its sources live\n", + " live.discard(dest)\n", + " if isinstance(src1, str): live.add(src1)\n", + " if isinstance(src2, str): live.add(src2)\n", + " new_program_reversed.append(instr)\n", + " else:\n", + " # Dead β€” replace with NOP\n", + " new_program_reversed.append({\n", + " \"op\": \"NOP\", \"dest\": None, \"src1\": None, \"src2\": None,\n", + " })\n", + "\n", + " else:\n", + " # Unknown op β€” defensive pass-through\n", + " new_program_reversed.append(instr)\n", + "\n", + " # Reverse back to forward order, then strip NOPs\n", + " new_program = list(reversed(new_program_reversed))\n", + " new_program = [instr for instr in new_program if instr[\"op\"] != \"NOP\"]\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "5vo-KgaffTFR" + }, + "execution_count": 8, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for dead_code_elimination ===\n", + "\n", + "# Test 1: simple dead variable\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None}, # never used\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result1 = dead_code_elimination(test1)\n", + "print(\"Test 1 (kill unused CONST):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: 'dead' instruction removed, 2 instructions remain\n", + "\n", + "# Test 2: chain of dead computation\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"x\", \"src1\": \"a\", \"src2\": \"b\"}, # x never used β†’ dead\n", + " {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 7, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result2 = dead_code_elimination(test2)\n", + "print(\"\\nTest 2 (kill unused ADD and its feeders... but only if feeders are also unused):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: x's ADD killed. a and b also killed (only fed into x, which is dead).\n", + "# Final: just CONST c=7, STORE.\n", + "\n", + "# Test 3: STORE always preserved\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result3 = dead_code_elimination(test3)\n", + "print(\"\\nTest 3 (STORE preserved, feeder kept live):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: both instructions unchanged\n", + "\n", + "# Test 4: variable used by STORE is live, even if defined far above\n", + "test4 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None}, # used by STORE β†’ live\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 9, \"src2\": None}, # never used β†’ dead\n", + " {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 1, \"src2\": None}, # never used β†’ dead\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result4 = dead_code_elimination(test4)\n", + "print(\"\\nTest 4 (only 'a' is live, b/c killed):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: a's CONST + STORE remain. b, c stripped.\n", + "\n", + "# Test 5: combine CF + DCE on a generated program\n", + "random.seed(42)\n", + "prog = generate_level_1()\n", + "print(\"\\nTest 5 (CF then DCE on seed=42):\")\n", + "print(\"Original:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"After CF:\")\n", + "for i, instr in enumerate(after_cf): print(f\" {i}: {instr}\")\n", + "\n", + "after_dce = dead_code_elimination(after_cf)\n", + "print(\"After CF + DCE:\")\n", + "for i, instr in enumerate(after_dce): print(f\" {i}: {instr}\")\n", + "# Expected: original 4 instructions become much shorter β€” dead 'v1' eliminated, v2 folded\n", + "\n", + "# Test 6: idempotence β€” running DCE twice gives same result\n", + "test6 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "once = dead_code_elimination(test6)\n", + "twice = dead_code_elimination(once)\n", + "print(\"\\nTest 6 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "vv8QRFRlvzdq", + "outputId": "cd7c4f43-81b4-41e8-b909-415ac97008a7" + }, + "execution_count": 9, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (kill unused CONST):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 2 (kill unused ADD and its feeders... but only if feeders are also unused):\n", + " 0: {'op': 'CONST', 'dest': 'c', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 3 (STORE preserved, feeder kept live):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 4 (only 'a' is live, b/c killed):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 5 (CF then DCE on seed=42):\n", + "Original:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After CF:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After CF + DCE:\n", + " 0: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "Test 6 (idempotence): PASS\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def peephole_optimization(program):\n", + " \"\"\"\n", + " Single-instruction peephole rewrites. Replaces expensive ops with\n", + " cheaper equivalents when operands hit special values (0, 1, 2, self).\n", + "\n", + " Patterns (all preserve semantics):\n", + " MUL by 2 -> ADD x + x\n", + " MUL by 1 -> CONST (the other operand)\n", + " MUL by 0 -> CONST 0\n", + " ADD with 0 -> CONST (the other operand)\n", + " SUB x - x -> CONST 0\n", + " DIV by 1 -> CONST (the dividend)\n", + "\n", + " Notes:\n", + " - Operates on individual instructions; no cross-instruction state.\n", + " - Relies on constant_folding having propagated literal values into operands.\n", + " - Returns fresh instruction dicts; never aliases inputs.\n", + "\n", + " Returns a new list. Never raises, never returns None.\n", + " \"\"\"\n", + " new_program = []\n", + "\n", + " for instr in program:\n", + " instr = instr.copy()\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + "\n", + " # === MUL patterns ===\n", + " if op == \"MUL\":\n", + " # MUL by 0 -> CONST 0\n", + " if src1 == 0 or src2 == 0:\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n", + " continue\n", + " # MUL by 1 -> CONST (other operand) if the other operand is a literal,\n", + " # otherwise leave alone (we don't want to introduce a useless\n", + " # \"CONST dest = some_var_name\" β€” that's not a valid CONST).\n", + " if src1 == 1 and isinstance(src2, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n", + " continue\n", + " if src2 == 1 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " # MUL by 2 -> ADD x + x (when one operand is literal 2 and the other is a var)\n", + " if src1 == 2 and isinstance(src2, str):\n", + " new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src2, \"src2\": src2})\n", + " continue\n", + " if src2 == 2 and isinstance(src1, str):\n", + " new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src1, \"src2\": src1})\n", + " continue\n", + " # No pattern matched β€” keep as is\n", + " new_program.append(instr)\n", + "\n", + " # === ADD patterns ===\n", + " elif op == \"ADD\":\n", + " # ADD with 0 -> CONST (other operand) if the other operand is a literal\n", + " if src1 == 0 and isinstance(src2, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n", + " continue\n", + " if src2 == 0 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " # If \"ADD x + 0\" where x is a variable, we'd want a copy β€” but our schema\n", + " # has no MOV/COPY op. Leave alone; CF + DCE handle the rest in practice.\n", + " new_program.append(instr)\n", + "\n", + " # === SUB patterns ===\n", + " elif op == \"SUB\":\n", + " # SUB x - x -> CONST 0 (only if both sources are the same string variable)\n", + " if isinstance(src1, str) and isinstance(src2, str) and src1 == src2:\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n", + " continue\n", + " # SUB x - 0 with literal 0 on src2 β€” we'd want a copy; skip (no COPY op)\n", + " new_program.append(instr)\n", + "\n", + " # === DIV patterns ===\n", + " elif op == \"DIV\":\n", + " # DIV by 1 -> CONST (dividend) if dividend is a literal\n", + " if src2 == 1 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " new_program.append(instr)\n", + "\n", + " else:\n", + " # CONST, LOAD, STORE, NOP β€” pass through\n", + " new_program.append(instr)\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "IaczDNeLv5PW" + }, + "execution_count": 10, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for peephole_optimization ===\n", + "\n", + "# Test 1: MUL by 2 -> ADD x+x\n", + "test1 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result1 = peephole_optimization(test1)\n", + "print(\"Test 1 (MUL by 2 -> ADD self):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes ADD y+y\n", + "\n", + "# Test 2: MUL by 0 -> CONST 0\n", + "test2 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 0},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result2 = peephole_optimization(test2)\n", + "print(\"\\nTest 2 (MUL by 0 -> CONST 0):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes CONST x = 0\n", + "\n", + "# Test 3: MUL by 1 (both literals)\n", + "test3 = [\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": 7, \"src2\": 1},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result3 = peephole_optimization(test3)\n", + "print(\"\\nTest 3 (MUL 7*1 -> CONST 7):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 0 becomes CONST x = 7\n", + "\n", + "# Test 4: SUB x - x -> CONST 0\n", + "test4 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"SUB\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result4 = peephole_optimization(test4)\n", + "print(\"\\nTest 4 (SUB y-y -> CONST 0):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes CONST x = 0\n", + "\n", + "# Test 5: ADD with 0\n", + "test5 = [\n", + " {\"op\": \"ADD\", \"dest\": \"x\", \"src1\": 0, \"src2\": 5},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result5 = peephole_optimization(test5)\n", + "print(\"\\nTest 5 (ADD 0+5 -> CONST 5):\")\n", + "for i, instr in enumerate(result5): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 0 becomes CONST x = 5\n", + "\n", + "# Test 6: full pipeline β€” CF then peephole on a generated program\n", + "random.seed(7)\n", + "prog = generate_level_1()\n", + "print(\"\\nTest 6 (CF then peephole on seed=7):\")\n", + "print(\"Original:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"After CF:\")\n", + "for i, instr in enumerate(after_cf): print(f\" {i}: {instr}\")\n", + "after_peep = peephole_optimization(after_cf)\n", + "print(\"After CF + peephole:\")\n", + "for i, instr in enumerate(after_peep): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 7: idempotence\n", + "test7 = [\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "once = peephole_optimization(test7)\n", + "twice = peephole_optimization(once)\n", + "print(\"\\nTest 7 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")\n", + "\n", + "# Test 8: no false fires β€” vanilla program shouldn't get rewritten\n", + "test8 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"}, # y*y, no peephole pattern\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result8 = peephole_optimization(test8)\n", + "print(\"\\nTest 8 (no false rewrite on y*y):\")\n", + "for i, instr in enumerate(result8): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 unchanged (still MUL y*y)" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "DWDUZUQvwXye", + "outputId": "88388440-df06-4f35-947f-44ba9b8bb833" + }, + "execution_count": 11, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (MUL by 2 -> ADD self):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'ADD', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 2 (MUL by 0 -> CONST 0):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 3 (MUL 7*1 -> CONST 7):\n", + " 0: {'op': 'CONST', 'dest': 'x', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 4 (SUB y-y -> CONST 0):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 5 (ADD 0+5 -> CONST 5):\n", + " 0: {'op': 'CONST', 'dest': 'x', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 6 (CF then peephole on seed=7):\n", + "Original:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "After CF:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "After CF + peephole:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "Test 7 (idempotence): PASS\n", + "\n", + "Test 8 (no false rewrite on y*y):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'MUL', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def dump_ir(program_data):\n", + " \"\"\"\n", + " Convert a TAC program into a readable string for the LLM observation.\n", + "\n", + " Format:\n", + " // OBSERVABLE OUT: mem[0], mem[1]\n", + " 0: v0 = 3 # CONST [1 cycle]\n", + " 1: v1 = 5 # CONST [1 cycle]\n", + " 2: v2 = v0 + v1 # ADD [1 cycle]\n", + " 3: mem[addr0] = v2 # STORE [4 cycles]\n", + " // TOTAL: 7 cycles\n", + "\n", + " Args:\n", + " program_data: dict with keys 'instructions' and 'observable_addrs',\n", + " OR a raw list of instructions (legacy).\n", + "\n", + " Returns:\n", + " A multi-line string suitable for inclusion in an LLM prompt.\n", + " \"\"\"\n", + " # Accept both forms β€” full program dict or raw instruction list\n", + " if isinstance(program_data, dict):\n", + " instructions = program_data[\"instructions\"]\n", + " observable_addrs = program_data.get(\"observable_addrs\", [])\n", + " else:\n", + " instructions = program_data\n", + " observable_addrs = []\n", + "\n", + " lines = []\n", + "\n", + " # Header: observable outputs\n", + " if observable_addrs:\n", + " addr_str = \", \".join(f\"mem[{a}]\" for a in observable_addrs)\n", + " lines.append(f\"// OBSERVABLE OUT: {addr_str}\")\n", + "\n", + " # Body: each instruction in human-readable form\n", + " total_cycles = 0\n", + " for i, instr in enumerate(instructions):\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + " cost = CYCLE_COSTS.get(op, 0)\n", + " total_cycles += cost\n", + "\n", + " # Render the instruction body\n", + " if op == \"CONST\":\n", + " body = f\"{dest} = {src1}\"\n", + " elif op == \"ADD\":\n", + " body = f\"{dest} = {src1} + {src2}\"\n", + " elif op == \"SUB\":\n", + " body = f\"{dest} = {src1} - {src2}\"\n", + " elif op == \"MUL\":\n", + " body = f\"{dest} = {src1} * {src2}\"\n", + " elif op == \"DIV\":\n", + " body = f\"{dest} = {src1} // {src2}\"\n", + " elif op == \"LOAD\":\n", + " body = f\"{dest} = mem[{src1}]\"\n", + " elif op == \"STORE\":\n", + " body = f\"mem[{dest}] = {src1}\"\n", + " elif op == \"NOP\":\n", + " body = \"nop\"\n", + " else:\n", + " body = f\"\"\n", + "\n", + " cost_label = f\"{cost} cycle\" if cost == 1 else f\"{cost} cycles\"\n", + " lines.append(f\"{i}: {body:<35} # {op:<6} [{cost_label}]\")\n", + "\n", + " # Footer: total cycles\n", + " lines.append(f\"// TOTAL: {total_cycles} cycles\")\n", + "\n", + " return \"\\n\".join(lines)" + ], + "metadata": { + "id": "yB_yMNOZwaPT" + }, + "execution_count": 12, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for state translator ===\n", + "\n", + "# Test 1: full generated program\n", + "random.seed(1)\n", + "prog = generate_level_1()\n", + "print(\"Test 1 (raw seed=1 program):\")\n", + "print(dump_ir(prog))\n", + "\n", + "# Test 2: after CF\n", + "print(\"\\nTest 2 (after CF):\")\n", + "folded = constant_folding(prog[\"instructions\"])\n", + "prog_after_cf = {**prog, \"instructions\": folded}\n", + "print(dump_ir(prog_after_cf))\n", + "\n", + "# Test 3: after CF + DCE β€” should show fewer instructions, lower total\n", + "print(\"\\nTest 3 (after CF + DCE):\")\n", + "optimized = dead_code_elimination(folded)\n", + "prog_optimized = {**prog, \"instructions\": optimized}\n", + "print(dump_ir(prog_optimized))\n", + "\n", + "# Test 4: every op type at least once\n", + "test4_instructions = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"b\", \"src1\": \"a\", \"src2\": 3},\n", + " {\"op\": \"SUB\", \"dest\": \"c\", \"src1\": \"b\", \"src2\": \"a\"},\n", + " {\"op\": \"MUL\", \"dest\": \"d\", \"src1\": \"c\", \"src2\": 2},\n", + " {\"op\": \"DIV\", \"dest\": \"e\", \"src1\": \"d\", \"src2\": 4},\n", + " {\"op\": \"LOAD\", \"dest\": \"f\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"e\", \"src2\": None},\n", + " {\"op\": \"NOP\", \"dest\": None, \"src1\": None, \"src2\": None},\n", + "]\n", + "test4_data = {\n", + " \"instructions\": test4_instructions,\n", + " \"observable_addrs\": [0, 1],\n", + "}\n", + "print(\"\\nTest 4 (every op type):\")\n", + "print(dump_ir(test4_data))\n", + "\n", + "# Test 5: empty program shouldn't crash\n", + "print(\"\\nTest 5 (empty program):\")\n", + "print(dump_ir({\"instructions\": [], \"observable_addrs\": [0]}))" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nWoFDPTH-Def", + "outputId": "69da36d2-569e-4175-d026-835bc09d9359" + }, + "execution_count": 13, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (raw seed=1 program):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 2 # CONST [1 cycle]\n", + "2: v2 = v1 + v1 # ADD [1 cycle]\n", + "3: v3 = v2 * v1 # MUL [3 cycles]\n", + "4: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 10 cycles\n", + "\n", + "Test 2 (after CF):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 2 # CONST [1 cycle]\n", + "2: v2 = 4 # CONST [1 cycle]\n", + "3: v3 = 8 # CONST [1 cycle]\n", + "4: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 8 cycles\n", + "\n", + "Test 3 (after CF + DCE):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v3 = 8 # CONST [1 cycle]\n", + "1: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n", + "\n", + "Test 4 (every op type):\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: a = 5 # CONST [1 cycle]\n", + "1: b = a + 3 # ADD [1 cycle]\n", + "2: c = b - a # SUB [1 cycle]\n", + "3: d = c * 2 # MUL [3 cycles]\n", + "4: e = d // 4 # DIV [5 cycles]\n", + "5: f = mem[addr0] # LOAD [4 cycles]\n", + "6: mem[addr0] = e # STORE [4 cycles]\n", + "7: nop # NOP [0 cycles]\n", + "// TOTAL: 19 cycles\n", + "\n", + "Test 5 (empty program):\n", + "// OBSERVABLE OUT: mem[0]\n", + "// TOTAL: 0 cycles\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_2():\n", + " \"\"\"\n", + " Generate a Level 2 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 8-12 instructions\n", + " - 3-5 CONSTs (some literal, some used in arithmetic)\n", + " - 3-5 arithmetic ops (ADD, SUB, MUL, DIV) with chaining\n", + " - 1-3 dead variables (DCE opportunities)\n", + " - 1 LOAD from initial memory whose result is GUARANTEED to flow into\n", + " the final STORE (forces agent to reason around unfoldable runtime values)\n", + " - 1 STORE at the end (observable output)\n", + " - DIV always uses literal non-zero divisor (Design Assumption #4)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 3-5 CONSTs\n", + " num_consts = random.randint(3, 5)\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + "\n", + " # Step 2: 1 LOAD from initial memory\n", + " load_addr_var = \"addr_in\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": load_addr_var, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + "\n", + " # Step 3: 3-5 arithmetic ops, chained\n", + " num_arith = random.randint(3, 5)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " src1 = random.choice(available_vars)\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4])\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3])\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 4: 1-3 extra dead CONSTs (DCE targets)\n", + " num_dead = random.randint(1, 3)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 20)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # === Step 5 (NEW): force LOAD-into-chain dependency ===\n", + " # Inject a final ADD that combines last_result with the loaded variable.\n", + " # This guarantees the LOAD result flows into the STORE β€” DCE can no longer\n", + " # eliminate it, and CF cannot collapse the entire program into a CONST.\n", + " final_dest = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": final_dest,\n", + " \"src1\": last_result,\n", + " \"src2\": loaded_var,\n", + " })\n", + " last_result = final_dest\n", + "\n", + " # Step 6: STORE the final result\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr_in\": 1,\n", + " }\n", + " initial_mem = {1: random.randint(1, 20)}\n", + "\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": [0],\n", + " }" + ], + "metadata": { + "id": "SUqa7odp3wB_" + }, + "execution_count": 14, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 2 generator ===\n", + "\n", + "# Test 1: spot-check a few seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_2()\n", + " print(f\"\\n=== Level 2, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + " print(f\"initial_vars: {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + "\n", + "# Test 2: full pipeline (CF + DCE + peephole) on a Level 2 program\n", + "print(\"\\n=== Full optimization pipeline on Level 2 (seed=42) ===\")\n", + "random.seed(42)\n", + "prog = generate_level_2()\n", + "print(\"ORIGINAL:\")\n", + "print(dump_ir(prog))\n", + "\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"\\nAFTER CF:\")\n", + "print(dump_ir({**prog, \"instructions\": after_cf}))\n", + "\n", + "after_dce = dead_code_elimination(after_cf)\n", + "print(\"\\nAFTER CF + DCE:\")\n", + "print(dump_ir({**prog, \"instructions\": after_dce}))\n", + "\n", + "after_peep = peephole_optimization(after_dce)\n", + "print(\"\\nAFTER CF + DCE + PEEPHOLE:\")\n", + "print(dump_ir({**prog, \"instructions\": after_peep}))" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "D0XjNetU5gIo", + "outputId": "dbcda9e6-ca12-42eb-9ab4-87dcfc15d157" + }, + "execution_count": 15, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 2, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 1}\n", + "\n", + "=== Level 2, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v7 = 19 # CONST [1 cycle]\n", + "2: v1 = 2 # CONST [1 cycle]\n", + "3: v2 = 5 # CONST [1 cycle]\n", + "4: v3 = mem[addr_in] # LOAD [4 cycles]\n", + "5: v4 = v3 // 2 # DIV [5 cycles]\n", + "6: v5 = v3 + v3 # ADD [1 cycle]\n", + "7: v6 = v2 + v5 # ADD [1 cycle]\n", + "8: v8 = v6 + v3 # ADD [1 cycle]\n", + "9: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 20 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 11}\n", + "\n", + "=== Level 2, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 3 # CONST [1 cycle]\n", + "1: v1 = 7 # CONST [1 cycle]\n", + "2: v2 = 1 # CONST [1 cycle]\n", + "3: v3 = 2 # CONST [1 cycle]\n", + "4: v4 = mem[addr_in] # LOAD [4 cycles]\n", + "5: v5 = v0 + v4 # ADD [1 cycle]\n", + "6: v6 = v3 - 3 # SUB [1 cycle]\n", + "7: v7 = v4 + 3 # ADD [1 cycle]\n", + "8: v10 = 2 # CONST [1 cycle]\n", + "9: v8 = v1 + v3 # ADD [1 cycle]\n", + "10: v9 = v6 + v0 # ADD [1 cycle]\n", + "11: v11 = v9 + v4 # ADD [1 cycle]\n", + "12: mem[addr0] = v11 # STORE [4 cycles]\n", + "// TOTAL: 19 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 5}\n", + "\n", + "=== Full optimization pipeline on Level 2 (seed=42) ===\n", + "ORIGINAL:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "\n", + "AFTER CF:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = 6 # CONST [1 cycle]\n", + "10: v7 = 1 # CONST [1 cycle]\n", + "11: v8 = -2 # CONST [1 cycle]\n", + "12: v12 = -2 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 20 cycles\n", + "\n", + "AFTER CF + DCE:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", + "\n", + "AFTER CF + DCE + PEEPHOLE:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Compose helper + evaluation harness ===\n", + "\n", + "# Registry mapping pass names (the LLM's action vocabulary) to functions.\n", + "PASS_REGISTRY = {\n", + " \"constant_folding\": constant_folding,\n", + " \"dead_code_elimination\": dead_code_elimination,\n", + " \"peephole_optimization\": peephole_optimization,\n", + "}\n", + "\n", + "\n", + "def count_cycles(instructions):\n", + " \"\"\"Sum cycle costs across an instruction list.\"\"\"\n", + " return sum(CYCLE_COSTS.get(instr[\"op\"], 0) for instr in instructions)\n", + "\n", + "\n", + "def apply_passes(instructions, sequence):\n", + " \"\"\"\n", + " Apply a sequence of passes (by name) to an instruction list.\n", + "\n", + " Args:\n", + " instructions: list of TAC instruction dicts.\n", + " sequence: list of pass-name strings, e.g. [\"constant_folding\", \"dead_code_elimination\"].\n", + "\n", + " Returns:\n", + " new instruction list (does not mutate input).\n", + "\n", + " Raises:\n", + " ValueError if an unknown pass name is provided. (This is intentional β€”\n", + " if Harshal's LLM emits a garbage pass name, we want to know loudly.)\n", + " \"\"\"\n", + " current = [instr.copy() for instr in instructions] # defensive copy\n", + " for pass_name in sequence:\n", + " if pass_name not in PASS_REGISTRY:\n", + " raise ValueError(f\"Unknown pass: {pass_name!r}. Known: {list(PASS_REGISTRY.keys())}\")\n", + " current = PASS_REGISTRY[pass_name](current)\n", + " return current\n", + "\n", + "\n", + "def evaluate(program, sequence):\n", + " \"\"\"\n", + " Run a pass sequence on a program and return cycle counts before/after.\n", + "\n", + " Args:\n", + " program: full program dict (with 'instructions', 'observable_addrs', etc.)\n", + " sequence: list of pass names.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - original_cycles: int\n", + " - optimized_cycles: int\n", + " - optimized_instructions: list[dict]\n", + " - cycle_reduction: int (original - optimized)\n", + " - reduction_pct: float (0.0 to 1.0)\n", + " \"\"\"\n", + " original_instructions = program[\"instructions\"]\n", + " original_cycles = count_cycles(original_instructions)\n", + "\n", + " optimized_instructions = apply_passes(original_instructions, sequence)\n", + " optimized_cycles = count_cycles(optimized_instructions)\n", + "\n", + " cycle_reduction = original_cycles - optimized_cycles\n", + " reduction_pct = cycle_reduction / original_cycles if original_cycles > 0 else 0.0\n", + "\n", + " return {\n", + " \"original_cycles\": original_cycles,\n", + " \"optimized_cycles\": optimized_cycles,\n", + " \"optimized_instructions\": optimized_instructions,\n", + " \"cycle_reduction\": cycle_reduction,\n", + " \"reduction_pct\": reduction_pct,\n", + " }" + ], + "metadata": { + "id": "Vo-9MitD5hwM" + }, + "execution_count": 17, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for compose helper + eval harness ===\n", + "\n", + "# Test 1: count_cycles on a known instruction list\n", + "test1_instrs = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"b\", \"src1\": \"a\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"b\", \"src2\": None},\n", + "]\n", + "print(f\"Test 1 (count_cycles): {count_cycles(test1_instrs)}\")\n", + "# Expected: 1 + 3 + 4 = 8\n", + "\n", + "# Test 2: apply_passes with single pass\n", + "result = apply_passes(test1_instrs, [\"constant_folding\"])\n", + "print(f\"\\nTest 2 (apply CF only):\")\n", + "for i, instr in enumerate(result): print(f\" {i}: {instr}\")\n", + "# Expected: MUL becomes ... actually 5*2 β†’ CONST 10. Wait β€” 'a' is const 5,\n", + "# but src2=2 is literal. Both resolve. CF should fold: b = CONST 10.\n", + "\n", + "# Test 3: apply_passes with sequence\n", + "result = apply_passes(test1_instrs, [\"constant_folding\", \"dead_code_elimination\"])\n", + "print(f\"\\nTest 3 (CF + DCE):\")\n", + "for i, instr in enumerate(result): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 4: full evaluate on a generated Level 2 program\n", + "random.seed(42)\n", + "prog = generate_level_2()\n", + "print(f\"\\nTest 4 (evaluate on Level 2 seed=42):\")\n", + "print(f\"Original program ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "eval_result = evaluate(prog, [\"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\"])\n", + "print(f\"\\nEvaluation result:\")\n", + "print(f\" original_cycles : {eval_result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {eval_result['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {eval_result['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {eval_result['reduction_pct']:.1%}\")\n", + "print(f\"\\nOptimized program:\")\n", + "print(dump_ir({**prog, \"instructions\": eval_result[\"optimized_instructions\"]}))\n", + "\n", + "# Test 5: unknown pass name raises clearly\n", + "try:\n", + " apply_passes(test1_instrs, [\"fake_pass\"])\n", + " print(\"\\nTest 5 (unknown pass): FAIL β€” should have raised\")\n", + "except ValueError as e:\n", + " print(f\"\\nTest 5 (unknown pass raises): PASS β€” {e}\")\n", + "\n", + "# Test 6: empty sequence is a no-op\n", + "result = apply_passes(test1_instrs, [])\n", + "matches = result == test1_instrs\n", + "print(f\"\\nTest 6 (empty sequence is no-op): {'PASS' if matches else 'FAIL'}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nNLyvg2JA-Xz", + "outputId": "57a62a25-330a-4c4a-cbac-83e84d64a14d" + }, + "execution_count": 18, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (count_cycles): 8\n", + "\n", + "Test 2 (apply CF only):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 10, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'b', 'src2': None}\n", + "\n", + "Test 3 (CF + DCE):\n", + " 0: {'op': 'CONST', 'dest': 'b', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'b', 'src2': None}\n", + "\n", + "Test 4 (evaluate on Level 2 seed=42):\n", + "Original program (24 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "\n", + "Evaluation result:\n", + " original_cycles : 24\n", + " optimized_cycles : 9\n", + " cycle_reduction : 15\n", + " reduction_pct : 62.5%\n", + "\n", + "Optimized program:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", + "\n", + "Test 5 (unknown pass raises): PASS β€” Unknown pass: 'fake_pass'. Known: ['constant_folding', 'dead_code_elimination', 'peephole_optimization']\n", + "\n", + "Test 6 (empty sequence is no-op): PASS\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Reverse Passes (Soham) ===\n", + "# Two deoptimization passes that temporarily complicate the IR.\n", + "# Rationale: a pure-simplification pass library can't escape local optima\n", + "# where the next forward step has nothing to bite on. Reverse passes create\n", + "# new structure (a fresh CONST + ADD pair, or a redundant binop) that gives\n", + "# subsequent forward passes a different starting point.\n", + "#\n", + "# Both passes:\n", + "# - return a NEW instruction list (do not mutate input)\n", + "# - use fresh variable names that don't collide with anything in the program\n", + "# - are deterministic (pick the FIRST eligible site)\n", + "# - never raise; if no eligible site exists, return an unchanged copy\n", + "\n", + "import re\n", + "\n", + "# Module-level counter for fresh variable name generation.\n", + "_FRESH_COUNTER = 0\n", + "\n", + "\n", + "def _collect_existing_names(program):\n", + " \"\"\"Return the set of all variable names appearing anywhere in the program.\"\"\"\n", + " names = set()\n", + " for ins in program:\n", + " for key in (\"dest\", \"src1\", \"src2\"):\n", + " v = ins.get(key)\n", + " if isinstance(v, str):\n", + " names.add(v)\n", + " return names\n", + "\n", + "\n", + "def _fresh_var(existing_names):\n", + " \"\"\"Return a name like '_rev0', '_rev1', ... not in existing_names.\"\"\"\n", + " global _FRESH_COUNTER\n", + " while True:\n", + " name = f\"_rev{_FRESH_COUNTER}\"\n", + " _FRESH_COUNTER += 1\n", + " if name not in existing_names:\n", + " return name\n", + "\n", + "\n", + "def _seed_counter_from_program(program):\n", + " \"\"\"\n", + " Advance the counter past any '_rev' already in the program, so the pass\n", + " is safe to apply to the output of a previous reverse pass without collision.\n", + " \"\"\"\n", + " global _FRESH_COUNTER\n", + " max_seen = -1\n", + " pat = re.compile(r\"^_rev(\\d+)$\")\n", + " for ins in program:\n", + " for key in (\"dest\", \"src1\", \"src2\"):\n", + " v = ins.get(key)\n", + " if isinstance(v, str):\n", + " m = pat.match(v)\n", + " if m:\n", + " max_seen = max(max_seen, int(m.group(1)))\n", + " if max_seen + 1 > _FRESH_COUNTER:\n", + " _FRESH_COUNTER = max_seen + 1\n", + "\n", + "\n", + "def expand_constant(program):\n", + " \"\"\"\n", + " Reverse pass: split a CONST V into (CONST a) + (ADD dest, _fresh, b)\n", + " where a = V // 2 and b = V - a, so a + b == V.\n", + "\n", + " Eligible: CONST whose src1 is an int with abs(src1) >= 2.\n", + " Picks the first eligible CONST in program order.\n", + " Returns a new list. If nothing is eligible, returns an unchanged copy.\n", + " \"\"\"\n", + " new_prog = [ins.copy() for ins in program]\n", + "\n", + " target_idx = None\n", + " for i, ins in enumerate(new_prog):\n", + " if ins[\"op\"] == \"CONST\" and isinstance(ins[\"src1\"], int) and abs(ins[\"src1\"]) >= 2:\n", + " target_idx = i\n", + " break\n", + "\n", + " if target_idx is None:\n", + " return new_prog\n", + "\n", + " target = new_prog[target_idx]\n", + " V = target[\"src1\"]\n", + " dest = target[\"dest\"]\n", + " a = V // 2\n", + " b = V - a # invariant: a + b == V, holds for negatives (Python floor div)\n", + "\n", + " _seed_counter_from_program(new_prog)\n", + " existing = _collect_existing_names(new_prog)\n", + " fresh = _fresh_var(existing)\n", + "\n", + " new_const = {\"op\": \"CONST\", \"dest\": fresh, \"src1\": a, \"src2\": None}\n", + " new_add = {\"op\": \"ADD\", \"dest\": dest, \"src1\": fresh, \"src2\": b}\n", + "\n", + " return new_prog[:target_idx] + [new_const, new_add] + new_prog[target_idx + 1:]\n", + "\n", + "\n", + "def duplicate_computation(program):\n", + " \"\"\"\n", + " Reverse pass: copy a binary op (ADD/SUB/MUL/DIV) with a fresh dest,\n", + " inserted immediately after the original. Creates a new live-range that\n", + " a forward pass can later reshape.\n", + "\n", + " Picks the first binary op in program order.\n", + " Returns a new list. If no binary op exists, returns an unchanged copy.\n", + " \"\"\"\n", + " BINARY_OPS = {\"ADD\", \"SUB\", \"MUL\", \"DIV\"}\n", + " new_prog = [ins.copy() for ins in program]\n", + "\n", + " target_idx = None\n", + " for i, ins in enumerate(new_prog):\n", + " if ins[\"op\"] in BINARY_OPS:\n", + " target_idx = i\n", + " break\n", + "\n", + " if target_idx is None:\n", + " return new_prog\n", + "\n", + " _seed_counter_from_program(new_prog)\n", + " existing = _collect_existing_names(new_prog)\n", + " fresh = _fresh_var(existing)\n", + "\n", + " original = new_prog[target_idx]\n", + " duplicate = original.copy()\n", + " duplicate[\"dest\"] = fresh\n", + "\n", + " return new_prog[:target_idx + 1] + [duplicate] + new_prog[target_idx + 1:]\n", + "\n", + "\n", + "# Register in the existing PASS_REGISTRY.\n", + "PASS_REGISTRY[\"expand_constant\"] = expand_constant\n", + "PASS_REGISTRY[\"duplicate_computation\"] = duplicate_computation\n", + "\n", + "print(f\"Registered {len(PASS_REGISTRY)} passes: {sorted(PASS_REGISTRY.keys())}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "saa_-mdHHDFg", + "outputId": "861c41eb-d519-4f72-f1be-20c40251700b" + }, + "execution_count": 19, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Registered 5 passes: ['constant_folding', 'dead_code_elimination', 'duplicate_computation', 'expand_constant', 'peephole_optimization']\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for reverse passes ===\n", + "\n", + "# Test 1: expand_constant on CONST 10 splits into 5 + 5\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"x\", \"src1\": 10, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result1 = expand_constant(test1)\n", + "assert len(result1) == 3\n", + "assert result1[0][\"op\"] == \"CONST\" and result1[0][\"src1\"] == 5\n", + "assert result1[1][\"op\"] == \"ADD\" and result1[1][\"dest\"] == \"x\" and result1[1][\"src2\"] == 5\n", + "assert result1[1][\"src1\"] == result1[0][\"dest\"]\n", + "assert result1[2] == test1[1]\n", + "print(f\"Test 1 (expand_constant on CONST 10): PASS β€” {result1[0]['dest']}=5, x={result1[0]['dest']}+5\")\n", + "\n", + "# Test 2: expand_constant skips CONST 1 and CONST 0; returns unchanged copy\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 1, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result2 = expand_constant(test2)\n", + "assert result2 == test2\n", + "assert result2 is not test2\n", + "print(\"Test 2 (expand_constant with no eligible CONST): PASS β€” returns unchanged copy\")\n", + "\n", + "# Test 3: duplicate_computation on ADD inserts a duplicate with fresh dest\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 4, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result3 = duplicate_computation(test3)\n", + "assert len(result3) == 5\n", + "assert result3[2][\"op\"] == \"ADD\" and result3[2][\"dest\"] == \"c\"\n", + "assert result3[3][\"op\"] == \"ADD\"\n", + "assert result3[3][\"dest\"] != \"c\"\n", + "assert result3[3][\"src1\"] == \"a\" and result3[3][\"src2\"] == \"b\"\n", + "assert result3[4] == test3[3]\n", + "print(f\"Test 3 (duplicate_computation on ADD): PASS β€” duplicate dest={result3[3]['dest']!r}\")\n", + "\n", + "# Test 4: length invariants and immutability\n", + "test4_in = [\n", + " {\"op\": \"CONST\", \"dest\": \"x\", \"src1\": 8, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"y\", \"src1\": \"x\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"y\", \"src2\": None},\n", + "]\n", + "e_out = expand_constant(test4_in)\n", + "d_out = duplicate_computation(test4_in)\n", + "assert len(e_out) == len(test4_in) + 1\n", + "assert len(d_out) == len(test4_in) + 1\n", + "assert test4_in == [\n", + " {\"op\": \"CONST\", \"dest\": \"x\", \"src1\": 8, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"y\", \"src1\": \"x\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"y\", \"src2\": None},\n", + "]\n", + "print(\"Test 4 (length invariants and immutability): PASS β€” expand +1, duplicate +1, input untouched\")\n", + "\n", + "# Test 5: negative constants split correctly\n", + "test5 = [{\"op\": \"CONST\", \"dest\": \"x\", \"src1\": -7, \"src2\": None}]\n", + "result5 = expand_constant(test5)\n", + "a, b = result5[0][\"src1\"], result5[1][\"src2\"]\n", + "assert a + b == -7\n", + "print(f\"Test 5 (expand_constant on negative): PASS β€” -7 = {a} + {b}\")\n", + "\n", + "# Test 6: registry has the new passes and they're callable through apply_passes\n", + "assert \"expand_constant\" in PASS_REGISTRY\n", + "assert \"duplicate_computation\" in PASS_REGISTRY\n", + "test6_prog = [\n", + " {\"op\": \"CONST\", \"dest\": \"x\", \"src1\": 6, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"y\", \"src1\": \"x\", \"src2\": 1},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"y\", \"src2\": None},\n", + "]\n", + "out6 = apply_passes(test6_prog, [\"expand_constant\", \"duplicate_computation\"])\n", + "assert len(out6) == len(test6_prog) + 2\n", + "print(f\"Test 6 (registered + callable via apply_passes): PASS β€” chained adds 2 instrs\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "NfTVatbrHGmf", + "outputId": "e511a95d-ec2e-4e09-bff1-cacfff98a05d" + }, + "execution_count": 20, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (expand_constant on CONST 10): PASS β€” _rev0=5, x=_rev0+5\n", + "Test 2 (expand_constant with no eligible CONST): PASS β€” returns unchanged copy\n", + "Test 3 (duplicate_computation on ADD): PASS β€” duplicate dest='_rev1'\n", + "Test 4 (length invariants and immutability): PASS β€” expand +1, duplicate +1, input untouched\n", + "Test 5 (expand_constant on negative): PASS β€” -7 = -4 + -3\n", + "Test 6 (registered + callable via apply_passes): PASS β€” chained adds 2 instrs\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Fixed-Point -O3 Baseline ===\n", + "\n", + "# The pass sequence applied in each iteration.\n", + "# Reasoning:\n", + "# CF first : propagates and folds constants, maximizing info for downstream passes\n", + "# PEEP next : rewrites special-value patterns (MUL by 2, ADD with 0) using propagated literals\n", + "# DCE last : sweeps up all newly-dead instructions\n", + "# Then loop until no further changes (fixed point).\n", + "O3_PASS_SEQUENCE = [\n", + " \"constant_folding\",\n", + " \"peephole_optimization\",\n", + " \"dead_code_elimination\",\n", + "]\n", + "\n", + "# Safety cap: prevent infinite loops on pathological inputs.\n", + "# In practice, fixed point is reached in 1-3 iterations.\n", + "MAX_FIXED_POINT_ITERATIONS = 10\n", + "\n", + "\n", + "def apply_o3_baseline(instructions):\n", + " \"\"\"\n", + " Apply the fixed-point -O3 baseline to an instruction list.\n", + "\n", + " Runs O3_PASS_SEQUENCE in a loop until the program stops changing\n", + " (fixed point reached) or MAX_FIXED_POINT_ITERATIONS is hit.\n", + "\n", + " Args:\n", + " instructions: list of TAC instruction dicts.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - optimized_instructions: list[dict] (final program)\n", + " - iterations: int (how many full sequence applications happened)\n", + " - converged: bool (True if fixed point reached, False if iter cap hit)\n", + " \"\"\"\n", + " current = [instr.copy() for instr in instructions]\n", + "\n", + " for iteration in range(1, MAX_FIXED_POINT_ITERATIONS + 1):\n", + " before = current\n", + " current = apply_passes(current, O3_PASS_SEQUENCE)\n", + " if current == before:\n", + " return {\n", + " \"optimized_instructions\": current,\n", + " \"iterations\": iteration,\n", + " \"converged\": True,\n", + " }\n", + "\n", + " return {\n", + " \"optimized_instructions\": current,\n", + " \"iterations\": MAX_FIXED_POINT_ITERATIONS,\n", + " \"converged\": False,\n", + " }\n", + "\n", + "\n", + "def evaluate_baseline(program):\n", + " \"\"\"\n", + " Run the -O3 baseline on a program. Convenience wrapper around apply_o3_baseline\n", + " that also computes cycle metrics.\n", + "\n", + " Args:\n", + " program: full program dict.\n", + "\n", + " Returns:\n", + " dict with original_cycles, optimized_cycles, cycle_reduction, reduction_pct,\n", + " iterations, converged, optimized_instructions.\n", + " \"\"\"\n", + " original_cycles = count_cycles(program[\"instructions\"])\n", + "\n", + " baseline_result = apply_o3_baseline(program[\"instructions\"])\n", + " optimized = baseline_result[\"optimized_instructions\"]\n", + " optimized_cycles = count_cycles(optimized)\n", + "\n", + " return {\n", + " \"original_cycles\": original_cycles,\n", + " \"optimized_cycles\": optimized_cycles,\n", + " \"cycle_reduction\": original_cycles - optimized_cycles,\n", + " \"reduction_pct\": (original_cycles - optimized_cycles) / original_cycles if original_cycles > 0 else 0.0,\n", + " \"iterations\": baseline_result[\"iterations\"],\n", + " \"converged\": baseline_result[\"converged\"],\n", + " \"optimized_instructions\": optimized,\n", + " }" + ], + "metadata": { + "id": "LFQfkN6hBCMS" + }, + "execution_count": 21, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for fixed-point -O3 baseline ===\n", + "\n", + "# Test 1: simple Level 1 program\n", + "random.seed(42)\n", + "prog1 = generate_level_1()\n", + "print(\"Test 1 (-O3 on Level 1, seed=42):\")\n", + "print(f\"Original ({count_cycles(prog1['instructions'])} cycles):\")\n", + "print(dump_ir(prog1))\n", + "\n", + "result1 = evaluate_baseline(prog1)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result1['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result1['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result1['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result1['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result1['iterations']}\")\n", + "print(f\" converged : {result1['converged']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog1, \"instructions\": result1[\"optimized_instructions\"]}))\n", + "\n", + "# Test 2: Level 2 β€” should keep LOAD alive\n", + "print(\"\\n\" + \"=\"*60)\n", + "random.seed(42)\n", + "prog2 = generate_level_2()\n", + "print(\"Test 2 (-O3 on Level 2, seed=42):\")\n", + "print(f\"Original ({count_cycles(prog2['instructions'])} cycles):\")\n", + "print(dump_ir(prog2))\n", + "\n", + "result2 = evaluate_baseline(prog2)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result2['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result2['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result2['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result2['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result2['iterations']}\")\n", + "print(f\" converged : {result2['converged']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog2, \"instructions\": result2[\"optimized_instructions\"]}))\n", + "\n", + "# Test 3: aggregate stats across many programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 2 programs):\")\n", + "total_original = 0\n", + "total_optimized = 0\n", + "iter_distribution = {}\n", + "for seed in range(100, 120):\n", + " random.seed(seed)\n", + " prog = generate_level_2()\n", + " result = evaluate_baseline(prog)\n", + " total_original += result[\"original_cycles\"]\n", + " total_optimized += result[\"optimized_cycles\"]\n", + " iter_distribution[result[\"iterations\"]] = iter_distribution.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_reduction_pct = (total_original - total_optimized) / total_original\n", + "print(f\" total original cycles : {total_original}\")\n", + "print(f\" total optimized cycles : {total_optimized}\")\n", + "print(f\" avg reduction pct : {avg_reduction_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_distribution}\")\n", + "\n", + "# Test 4: idempotence β€” running baseline twice gives same result\n", + "print(\"\\nTest 4 (idempotence β€” applying baseline to its own output is no-op):\")\n", + "random.seed(7)\n", + "prog4 = generate_level_2()\n", + "once = apply_o3_baseline(prog4[\"instructions\"])\n", + "twice = apply_o3_baseline(once[\"optimized_instructions\"])\n", + "matches = once[\"optimized_instructions\"] == twice[\"optimized_instructions\"]\n", + "print(f\" PASS: baseline output is at fixed point\" if matches else \" FAIL: applying baseline twice changed program\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "1cGy1HTSCccB", + "outputId": "f7f246bb-5ce3-442f-8a97-0d3f8c9f6ffe" + }, + "execution_count": 22, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (-O3 on Level 1, seed=42):\n", + "Original (7 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 1 # CONST [1 cycle]\n", + "1: v1 = 5 # CONST [1 cycle]\n", + "2: v2 = v0 + v0 # ADD [1 cycle]\n", + "3: mem[addr0] = v2 # STORE [4 cycles]\n", + "// TOTAL: 7 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 7\n", + " optimized_cycles : 5\n", + " cycle_reduction : 2\n", + " reduction_pct : 28.6%\n", + " iterations : 2\n", + " converged : True\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v2 = 2 # CONST [1 cycle]\n", + "1: mem[addr0] = v2 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n", + "\n", + "============================================================\n", + "Test 2 (-O3 on Level 2, seed=42):\n", + "Original (24 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 24\n", + " optimized_cycles : 9\n", + " cycle_reduction : 15\n", + " reduction_pct : 62.5%\n", + " iterations : 2\n", + " converged : True\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 2 programs):\n", + " total original cycles : 522\n", + " total optimized cycles : 231\n", + " avg reduction pct : 55.7%\n", + " iteration distribution : {2: 19, 3: 1}\n", + "\n", + "Test 4 (idempotence β€” applying baseline to its own output is no-op):\n", + " PASS: baseline output is at fixed point\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Held-Out Test Set ===\n", + "\n", + "# Seed partitioning convention (LOCKED β€” must match Harshal's training sampler):\n", + "TRAINING_SEED_MIN = 0\n", + "TRAINING_SEED_MAX = 9999\n", + "HOLDOUT_SEED_MIN = 10000\n", + "HOLDOUT_SEED_MAX = 19999\n", + "\n", + "# Map level number -> generator function.\n", + "# Levels 3-5 will be added as we build them.\n", + "LEVEL_GENERATORS = {\n", + " 1: generate_level_1,\n", + " 2: generate_level_2,\n", + "}\n", + "\n", + "\n", + "def generate_test_set(n_per_level=20, levels=None):\n", + " \"\"\"\n", + " Generate a deterministic held-out test set across difficulty levels.\n", + "\n", + " Uses seeds from the held-out range (10000-19999) only β€” guaranteed\n", + " not to overlap with training data (0-9999).\n", + "\n", + " Args:\n", + " n_per_level: number of programs per difficulty level (default 20).\n", + " levels: list of level numbers to include. Defaults to all available.\n", + "\n", + " Returns:\n", + " list of dicts, each with:\n", + " - level: int (which difficulty level it came from)\n", + " - seed: int (the seed used to generate it)\n", + " - program: full program dict\n", + " \"\"\"\n", + " if levels is None:\n", + " levels = sorted(LEVEL_GENERATORS.keys())\n", + "\n", + " test_set = []\n", + "\n", + " for level in levels:\n", + " if level not in LEVEL_GENERATORS:\n", + " raise ValueError(f\"No generator for level {level}. Available: {list(LEVEL_GENERATORS.keys())}\")\n", + "\n", + " generator = LEVEL_GENERATORS[level]\n", + "\n", + " # Use held-out seeds, offset by level so seeds don't repeat across levels.\n", + " # Level 1 uses 10000-10000+n, Level 2 uses 11000-11000+n, etc.\n", + " base_seed = HOLDOUT_SEED_MIN + (level - 1) * 1000\n", + "\n", + " for i in range(n_per_level):\n", + " seed = base_seed + i\n", + " random.seed(seed)\n", + " program = generator()\n", + " test_set.append({\n", + " \"level\": level,\n", + " \"seed\": seed,\n", + " \"program\": program,\n", + " })\n", + "\n", + " return test_set\n", + "\n", + "\n", + "def run_baseline_on_test_set(test_set):\n", + " \"\"\"\n", + " Run the fixed-point -O3 baseline on every program in the test set.\n", + "\n", + " Args:\n", + " test_set: list of {level, seed, program} dicts.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - per_program: list of per-program results (cycles before/after, level, seed)\n", + " - by_level: dict mapping level -> aggregate stats\n", + " - overall: aggregate stats across all programs\n", + " \"\"\"\n", + " per_program = []\n", + "\n", + " for entry in test_set:\n", + " result = evaluate_baseline(entry[\"program\"])\n", + " per_program.append({\n", + " \"level\": entry[\"level\"],\n", + " \"seed\": entry[\"seed\"],\n", + " \"original_cycles\": result[\"original_cycles\"],\n", + " \"optimized_cycles\": result[\"optimized_cycles\"],\n", + " \"cycle_reduction\": result[\"cycle_reduction\"],\n", + " \"reduction_pct\": result[\"reduction_pct\"],\n", + " \"iterations\": result[\"iterations\"],\n", + " })\n", + "\n", + " # Aggregate by level\n", + " by_level = {}\n", + " for entry in per_program:\n", + " lvl = entry[\"level\"]\n", + " if lvl not in by_level:\n", + " by_level[lvl] = {\n", + " \"n_programs\": 0,\n", + " \"total_original\": 0,\n", + " \"total_optimized\": 0,\n", + " \"sum_reduction_pct\": 0.0,\n", + " }\n", + " by_level[lvl][\"n_programs\"] += 1\n", + " by_level[lvl][\"total_original\"] += entry[\"original_cycles\"]\n", + " by_level[lvl][\"total_optimized\"] += entry[\"optimized_cycles\"]\n", + " by_level[lvl][\"sum_reduction_pct\"] += entry[\"reduction_pct\"]\n", + "\n", + " # Compute averages\n", + " for lvl, stats in by_level.items():\n", + " stats[\"avg_original\"] = stats[\"total_original\"] / stats[\"n_programs\"]\n", + " stats[\"avg_optimized\"] = stats[\"total_optimized\"] / stats[\"n_programs\"]\n", + " stats[\"avg_reduction_pct\"] = stats[\"sum_reduction_pct\"] / stats[\"n_programs\"]\n", + " # Aggregate reduction percent (different from average of percents)\n", + " stats[\"aggregate_reduction_pct\"] = (\n", + " (stats[\"total_original\"] - stats[\"total_optimized\"]) / stats[\"total_original\"]\n", + " if stats[\"total_original\"] > 0 else 0.0\n", + " )\n", + "\n", + " # Overall\n", + " total_orig = sum(e[\"original_cycles\"] for e in per_program)\n", + " total_opt = sum(e[\"optimized_cycles\"] for e in per_program)\n", + " overall = {\n", + " \"n_programs\": len(per_program),\n", + " \"total_original\": total_orig,\n", + " \"total_optimized\": total_opt,\n", + " \"aggregate_reduction_pct\": (total_orig - total_opt) / total_orig if total_orig > 0 else 0.0,\n", + " \"avg_reduction_pct\": sum(e[\"reduction_pct\"] for e in per_program) / len(per_program) if per_program else 0.0,\n", + " }\n", + "\n", + " return {\n", + " \"per_program\": per_program,\n", + " \"by_level\": by_level,\n", + " \"overall\": overall,\n", + " }" + ], + "metadata": { + "id": "6LuK5ln0Cd4K" + }, + "execution_count": 23, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for held-out test set ===\n", + "\n", + "# Test 1: deterministic β€” same call returns same programs\n", + "ts1 = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "ts2 = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "\n", + "# Compare instruction lists across the two calls\n", + "all_match = all(\n", + " a[\"program\"][\"instructions\"] == b[\"program\"][\"instructions\"]\n", + " for a, b in zip(ts1, ts2)\n", + ")\n", + "print(f\"Test 1 (deterministic test set): {'PASS' if all_match else 'FAIL'}\")\n", + "\n", + "# Test 2: seeds are in held-out range\n", + "ts = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "all_holdout = all(HOLDOUT_SEED_MIN <= e[\"seed\"] <= HOLDOUT_SEED_MAX for e in ts)\n", + "print(f\"Test 2 (all seeds in held-out range): {'PASS' if all_holdout else 'FAIL'}\")\n", + "\n", + "# Test 3: distribution across levels\n", + "print(f\"\\nTest 3 (test set composition for n_per_level=5, levels=[1,2]):\")\n", + "level_counts = {}\n", + "for e in ts: level_counts[e[\"level\"]] = level_counts.get(e[\"level\"], 0) + 1\n", + "print(f\" level counts: {level_counts}\")\n", + "print(f\" total: {len(ts)}\")\n", + "\n", + "# Test 4: full baseline run on a real test set (20 per level)\n", + "print(f\"\\n\" + \"=\"*60)\n", + "print(f\"Test 4 (baseline on full held-out test set, 20/level x 2 levels = 40 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "print(f\"\\n === BY LEVEL ===\")\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "print(f\"\\n === OVERALL ===\")\n", + "o = results[\"overall\"]\n", + "print(f\" n_programs : {o['n_programs']}\")\n", + "print(f\" total original cycles : {o['total_original']}\")\n", + "print(f\" total optimized cycles : {o['total_optimized']}\")\n", + "print(f\" aggregate reduction pct : {o['aggregate_reduction_pct']:.1%}\")\n", + "print(f\" avg per-program redux : {o['avg_reduction_pct']:.1%}\")\n", + "\n", + "# Test 5: per-program detail (first 3 entries) β€” useful for debugging\n", + "print(f\"\\n === PER-PROGRAM (first 3) ===\")\n", + "for entry in results[\"per_program\"][:3]:\n", + " print(f\" Level {entry['level']}, seed {entry['seed']}: \"\n", + " f\"{entry['original_cycles']} -> {entry['optimized_cycles']} \"\n", + " f\"({entry['reduction_pct']:.1%}, {entry['iterations']} iter)\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "bHmrTiFsFMsR", + "outputId": "bc812930-657e-41c6-fb34-4ceaa7483eb6" + }, + "execution_count": 24, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (deterministic test set): PASS\n", + "Test 2 (all seeds in held-out range): PASS\n", + "\n", + "Test 3 (test set composition for n_per_level=5, levels=[1,2]):\n", + " level counts: {1: 5, 2: 5}\n", + " total: 10\n", + "\n", + "============================================================\n", + "Test 4 (baseline on full held-out test set, 20/level x 2 levels = 40 programs):\n", + "\n", + " === BY LEVEL ===\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + "\n", + " === OVERALL ===\n", + " n_programs : 40\n", + " total original cycles : 703\n", + " total optimized cycles : 329\n", + " aggregate reduction pct : 53.2%\n", + " avg per-program redux : 50.8%\n", + "\n", + " === PER-PROGRAM (first 3) ===\n", + " Level 1, seed 10000: 11 -> 5 (54.5%, 2 iter)\n", + " Level 1, seed 10001: 10 -> 5 (50.0%, 2 iter)\n", + " Level 1, seed 10002: 10 -> 5 (50.0%, 2 iter)\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_3():\n", + " \"\"\"\n", + " Generate a Level 3 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 15-20 instructions\n", + " - 5-7 CONSTs\n", + " - 2 LOADs (both forced to flow into final STORE)\n", + " - 6-9 arithmetic ops with deeper chaining\n", + " - 2-4 dead variables sprinkled throughout\n", + " - 1 main STORE (optionally a secondary STORE)\n", + " - DIV always uses literal non-zero divisor (Design Assumption #4)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 5-7 CONSTs\n", + " num_consts = random.randint(5, 7)\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 15)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + "\n", + " # Step 2: 2 LOADs from initial memory\n", + " loaded_vars = []\n", + " for i in range(2):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: 6-9 arithmetic ops, chained\n", + " num_arith = random.randint(6, 9)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " src1 = random.choice(available_vars)\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3])\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 4: 2-4 dead CONSTs scattered throughout\n", + " num_dead = random.randint(2, 4)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 25)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 5: Force BOTH LOADed variables into the final result chain.\n", + " # Combine last_result with both loaded values via two ADDs.\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": after_load1,\n", + " \"src1\": last_result,\n", + " \"src2\": loaded_vars[0],\n", + " })\n", + "\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": after_load2,\n", + " \"src1\": after_load1,\n", + " \"src2\": loaded_vars[1],\n", + " })\n", + " last_result = after_load2\n", + "\n", + " # Step 6: Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " # Step 7 (optional, 40% chance): Secondary STORE to a different address.\n", + " # Stores an intermediate constant β€” gives the agent another observable to preserve.\n", + " has_secondary_store = random.random() < 0.4\n", + " observable_addrs = [0]\n", + " if has_secondary_store:\n", + " # Pick one of the early-defined CONSTs as the secondary value\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr1\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(1)\n", + "\n", + " # Step 8: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr_in0\": 2,\n", + " \"addr_in1\": 3,\n", + " }\n", + " initial_mem = {\n", + " 2: random.randint(1, 25),\n", + " 3: random.randint(1, 25),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 3 with the test set generator\n", + "LEVEL_GENERATORS[3] = generate_level_3" + ], + "metadata": { + "id": "LtziAqGJFPN4" + }, + "execution_count": 25, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 3 generator ===\n", + "\n", + "# Test 1: spot-check three seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_3()\n", + " print(f\"\\n=== Level 3, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + " print(f\"observable_addrs : {prog['observable_addrs']}\")\n", + " print(f\"initial_vars : {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + "\n", + "# Test 2: Full pipeline on Level 3, seed=42\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (-O3 baseline on Level 3, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_3()\n", + "print(f\"Original ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "result = evaluate_baseline(prog)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog, \"instructions\": result[\"optimized_instructions\"]}))\n", + "\n", + "# Test 3: aggregate stats across 20 Level 3 programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 3 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(200, 220):\n", + " random.seed(seed)\n", + " prog = generate_level_3()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 4: held-out test set with Level 3 included\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (held-out test set, Levels 1-3, 20 each):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n Overall: n={o['n_programs']}, \"\n", + " f\"agg reduction={o['aggregate_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "7Edy1LM1TJzg", + "outputId": "3e1adc51-f7d9-404e-ade8-a873ef5f1808" + }, + "execution_count": 26, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 3, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v4 = 4 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 3 # CONST [1 cycle]\n", + "7: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v9 = v0 + v0 # ADD [1 cycle]\n", + "10: v10 = v8 + 0 # ADD [1 cycle]\n", + "11: v16 = 4 # CONST [1 cycle]\n", + "12: v11 = v8 - v6 # SUB [1 cycle]\n", + "13: v15 = 3 # CONST [1 cycle]\n", + "14: v12 = v4 - v0 # SUB [1 cycle]\n", + "15: v13 = v5 - v4 # SUB [1 cycle]\n", + "16: v14 = v12 - 2 # SUB [1 cycle]\n", + "17: v17 = v14 + v7 # ADD [1 cycle]\n", + "18: v18 = v17 + v8 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 29 cycles\n", + "observable_addrs : [0]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 20, 3: 9}\n", + "\n", + "=== Level 3, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v15 = 10 # CONST [1 cycle]\n", + "1: v0 = 10 # CONST [1 cycle]\n", + "2: v1 = 14 # CONST [1 cycle]\n", + "3: v16 = 14 # CONST [1 cycle]\n", + "4: v2 = 13 # CONST [1 cycle]\n", + "5: v3 = 13 # CONST [1 cycle]\n", + "6: v4 = 2 # CONST [1 cycle]\n", + "7: v5 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v6 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v17 = 6 # CONST [1 cycle]\n", + "10: v7 = v3 + v3 # ADD [1 cycle]\n", + "11: v8 = v1 // 1 # DIV [5 cycles]\n", + "12: v9 = v0 // 3 # DIV [5 cycles]\n", + "13: v10 = v1 - v5 # SUB [1 cycle]\n", + "14: v11 = v10 + 0 # ADD [1 cycle]\n", + "15: v12 = v6 // 1 # DIV [5 cycles]\n", + "16: v13 = v7 - v8 # SUB [1 cycle]\n", + "17: v14 = v10 - v3 # SUB [1 cycle]\n", + "18: v18 = v14 + v5 # ADD [1 cycle]\n", + "19: v19 = v18 + v6 # ADD [1 cycle]\n", + "20: mem[addr0] = v19 # STORE [4 cycles]\n", + "21: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 46 cycles\n", + "observable_addrs : [0, 1]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 24, 3: 23}\n", + "\n", + "=== Level 3, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 3 # CONST [1 cycle]\n", + "1: v1 = 7 # CONST [1 cycle]\n", + "2: v2 = 11 # CONST [1 cycle]\n", + "3: v3 = 1 # CONST [1 cycle]\n", + "4: v4 = 2 # CONST [1 cycle]\n", + "5: v16 = 18 # CONST [1 cycle]\n", + "6: v5 = 14 # CONST [1 cycle]\n", + "7: v6 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v7 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v14 = 5 # CONST [1 cycle]\n", + "10: v15 = 4 # CONST [1 cycle]\n", + "11: v8 = v3 * v0 # MUL [3 cycles]\n", + "12: v9 = v1 + v3 # ADD [1 cycle]\n", + "13: v10 = v0 + v9 # ADD [1 cycle]\n", + "14: v11 = v10 + v10 # ADD [1 cycle]\n", + "15: v12 = v6 + v0 # ADD [1 cycle]\n", + "16: v13 = v2 - 2 # SUB [1 cycle]\n", + "17: v17 = v13 + v6 # ADD [1 cycle]\n", + "18: v18 = v17 + v7 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "20: mem[addr1] = v4 # STORE [4 cycles]\n", + "// TOTAL: 35 cycles\n", + "observable_addrs : [0, 1]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 21, 3: 7}\n", + "\n", + "============================================================\n", + "Test 2 (-O3 baseline on Level 3, seed=42):\n", + "Original (29 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v4 = 4 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 3 # CONST [1 cycle]\n", + "7: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v9 = v0 + v0 # ADD [1 cycle]\n", + "10: v10 = v8 + 0 # ADD [1 cycle]\n", + "11: v16 = 4 # CONST [1 cycle]\n", + "12: v11 = v8 - v6 # SUB [1 cycle]\n", + "13: v15 = 3 # CONST [1 cycle]\n", + "14: v12 = v4 - v0 # SUB [1 cycle]\n", + "15: v13 = v5 - v4 # SUB [1 cycle]\n", + "16: v14 = v12 - 2 # SUB [1 cycle]\n", + "17: v17 = v14 + v7 # ADD [1 cycle]\n", + "18: v18 = v17 + v8 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 29 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 29\n", + " optimized_cycles : 14\n", + " cycle_reduction : 15\n", + " reduction_pct : 51.7%\n", + " iterations : 2\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "1: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "2: v17 = 0 + v7 # ADD [1 cycle]\n", + "3: v18 = v17 + v8 # ADD [1 cycle]\n", + "4: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 14 cycles\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 3 programs):\n", + " total original cycles : 839\n", + " total optimized cycles : 329\n", + " avg reduction pct : 60.8%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 4 (held-out test set, Levels 1-3, 20 each):\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig=43.4, avg opt=17.6, agg reduction=59.5%\n", + "\n", + " Overall: n=60, agg reduction=56.7%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_4():\n", + " \"\"\"\n", + " Generate a Level 4 Toy-IR program with fan-out topology.\n", + "\n", + " Distinct from Levels 1-3: 2-3 'hub' variables fan out to drive\n", + " most arithmetic operations. A single optimization at a hub cascades\n", + " through many dependents, giving the agent measurable leverage from\n", + " correctly identifying high-value optimization targets.\n", + "\n", + " Characteristics:\n", + " - 20-28 instructions\n", + " - 5-7 CONSTs (some become hubs)\n", + " - 2 LOADs (one usually a hub)\n", + " - 10-14 arithmetic ops, ~70% using a hub variable as src1\n", + " - 3-5 dead variables\n", + " - 1-2 STOREs (both LOADs flow into final result)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 5-7 CONSTs\n", + " num_consts = random.randint(5, 7)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 15)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: 2 LOADs\n", + " loaded_vars = []\n", + " for i in range(2):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: Designate 2-3 hub variables.\n", + " # Hubs come from a mix of constants and loads to ensure CF can exploit\n", + " # constant hubs while non-constant (LOAD) hubs force the agent to reason\n", + " # about partially-resolvable structure.\n", + " num_hubs = random.randint(2, 3)\n", + " candidate_hubs = const_vars[:3] + loaded_vars # bias toward early-defined vars\n", + " hub_vars = random.sample(candidate_hubs, min(num_hubs, len(candidate_hubs)))\n", + "\n", + " # Step 4: 10-14 arithmetic ops, ~70% using a hub as src1\n", + " num_arith = random.randint(10, 14)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " # 70% chance src1 is a hub (drives the fan-out structure)\n", + " if random.random() < 0.70:\n", + " src1 = random.choice(hub_vars)\n", + " else:\n", + " src1 = random.choice(available_vars)\n", + "\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3])\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 5: 3-5 dead CONSTs scattered throughout\n", + " num_dead = random.randint(3, 5)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 30)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 6: Force both LOADed values into the final chain\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load1,\n", + " \"src1\": last_result, \"src2\": loaded_vars[0],\n", + " })\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load2,\n", + " \"src1\": after_load1, \"src2\": loaded_vars[1],\n", + " })\n", + " last_result = after_load2\n", + "\n", + " # Step 7: Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " # Step 8 (40% chance): secondary STORE\n", + " has_secondary_store = random.random() < 0.4\n", + " observable_addrs = [0]\n", + " if has_secondary_store:\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr1\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(1)\n", + "\n", + " # Step 9: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr_in0\": 2,\n", + " \"addr_in1\": 3,\n", + " }\n", + " initial_mem = {\n", + " 2: random.randint(1, 30),\n", + " 3: random.randint(1, 30),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 4\n", + "LEVEL_GENERATORS[4] = generate_level_4" + ], + "metadata": { + "id": "cyi06K_mVhUf" + }, + "execution_count": 27, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 4 generator ===\n", + "\n", + "# Test 1: spot-check three seeds β€” verify fan-out is visible\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_4()\n", + " print(f\"\\n=== Level 4, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + "\n", + "# Test 2: structural verification β€” count fan-out\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (verify fan-out structure):\")\n", + "random.seed(42)\n", + "prog = generate_level_4()\n", + "\n", + "# Count how often each variable is used as src1 or src2\n", + "usage_count = {}\n", + "for instr in prog[\"instructions\"]:\n", + " for src_field in (\"src1\", \"src2\"):\n", + " src = instr[src_field]\n", + " if isinstance(src, str):\n", + " usage_count[src] = usage_count.get(src, 0) + 1\n", + "\n", + "# Sort by usage\n", + "sorted_usage = sorted(usage_count.items(), key=lambda x: -x[1])\n", + "print(f\" Top 5 most-used variables (the 'hubs'):\")\n", + "for var, count in sorted_usage[:5]:\n", + " print(f\" {var}: used {count} times\")\n", + "print(f\" (In linear Level 3, top variables typically used 1-3 times.)\")\n", + "\n", + "# Test 3: baseline on Level 4\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 baseline on Level 4, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_4()\n", + "print(f\"Original ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "result = evaluate_baseline(prog)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "\n", + "# Test 4: aggregate stats on Level 4\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (-O3 across 20 random Level 4 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(300, 320):\n", + " random.seed(seed)\n", + " prog = generate_level_4()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 5: full held-out test set (Levels 1-4)\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 5 (held-out test set, Levels 1-4, 20 each = 80 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3, 4])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n Overall: n={o['n_programs']}, agg reduction={o['aggregate_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "m8yNd1VYVhuL", + "outputId": "74d1101d-ccdc-40ef-d68f-e94cd02fcf97" + }, + "execution_count": 28, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 4, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v25 = 15 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = 4 # CONST [1 cycle]\n", + "7: v6 = 3 # CONST [1 cycle]\n", + "8: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "10: v9 = v8 // 5 # DIV [5 cycles]\n", + "11: v10 = v8 + v6 # ADD [1 cycle]\n", + "12: v11 = v8 - v2 # SUB [1 cycle]\n", + "13: v12 = v0 // 1 # DIV [5 cycles]\n", + "14: v13 = v9 // 3 # DIV [5 cycles]\n", + "15: v14 = v0 + v1 # ADD [1 cycle]\n", + "16: v23 = 21 # CONST [1 cycle]\n", + "17: v15 = v0 * v9 # MUL [3 cycles]\n", + "18: v16 = v8 - v9 # SUB [1 cycle]\n", + "19: v17 = v12 + v8 # ADD [1 cycle]\n", + "20: v18 = v0 // 2 # DIV [5 cycles]\n", + "21: v24 = 13 # CONST [1 cycle]\n", + "22: v19 = v8 * v5 # MUL [3 cycles]\n", + "23: v20 = v8 - 2 # SUB [1 cycle]\n", + "24: v21 = v0 + 3 # ADD [1 cycle]\n", + "25: v22 = v18 * 2 # MUL [3 cycles]\n", + "26: v26 = v22 + v7 # ADD [1 cycle]\n", + "27: v27 = v26 + v8 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 64 cycles\n", + "\n", + "=== Level 4, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 14 # CONST [1 cycle]\n", + "2: v2 = 13 # CONST [1 cycle]\n", + "3: v3 = 13 # CONST [1 cycle]\n", + "4: v4 = 2 # CONST [1 cycle]\n", + "5: v21 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in0] # LOAD [4 cycles]\n", + "7: v6 = mem[addr_in1] # LOAD [4 cycles]\n", + "8: v7 = v0 // 4 # DIV [5 cycles]\n", + "9: v8 = v7 // 2 # DIV [5 cycles]\n", + "10: v9 = v0 + v8 # ADD [1 cycle]\n", + "11: v10 = v0 + v6 # ADD [1 cycle]\n", + "12: v22 = 17 # CONST [1 cycle]\n", + "13: v11 = v7 + v8 # ADD [1 cycle]\n", + "14: v12 = v5 - v4 # SUB [1 cycle]\n", + "15: v23 = 12 # CONST [1 cycle]\n", + "16: v13 = v10 + v1 # ADD [1 cycle]\n", + "17: v14 = v4 - v1 # SUB [1 cycle]\n", + "18: v15 = v8 * v14 # MUL [3 cycles]\n", + "19: v20 = 22 # CONST [1 cycle]\n", + "20: v16 = v6 // 3 # DIV [5 cycles]\n", + "21: v17 = v1 // 4 # DIV [5 cycles]\n", + "22: v18 = v1 - v5 # SUB [1 cycle]\n", + "23: v19 = v11 * v2 # MUL [3 cycles]\n", + "24: v24 = v19 + v5 # ADD [1 cycle]\n", + "25: v25 = v24 + v6 # ADD [1 cycle]\n", + "26: mem[addr0] = v25 # STORE [4 cycles]\n", + "// TOTAL: 56 cycles\n", + "\n", + "=== Level 4, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v24 = 12 # CONST [1 cycle]\n", + "1: v0 = 3 # CONST [1 cycle]\n", + "2: v1 = 7 # CONST [1 cycle]\n", + "3: v2 = 11 # CONST [1 cycle]\n", + "4: v3 = 1 # CONST [1 cycle]\n", + "5: v4 = 2 # CONST [1 cycle]\n", + "6: v5 = 14 # CONST [1 cycle]\n", + "7: v6 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v7 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v8 = v2 - 1 # SUB [1 cycle]\n", + "10: v9 = v2 + v3 # ADD [1 cycle]\n", + "11: v25 = 15 # CONST [1 cycle]\n", + "12: v10 = v2 + v0 # ADD [1 cycle]\n", + "13: v11 = v2 - 2 # SUB [1 cycle]\n", + "14: v22 = 23 # CONST [1 cycle]\n", + "15: v12 = v2 - 2 # SUB [1 cycle]\n", + "16: v13 = v2 + v9 # ADD [1 cycle]\n", + "17: v14 = v0 - v7 # SUB [1 cycle]\n", + "18: v15 = v2 // 2 # DIV [5 cycles]\n", + "19: v16 = v0 + v14 # ADD [1 cycle]\n", + "20: v17 = v0 * v5 # MUL [3 cycles]\n", + "21: v18 = v2 * 0 # MUL [3 cycles]\n", + "22: v19 = v0 * v18 # MUL [3 cycles]\n", + "23: v23 = 29 # CONST [1 cycle]\n", + "24: v20 = v0 // 4 # DIV [5 cycles]\n", + "25: v21 = v20 + 3 # ADD [1 cycle]\n", + "26: v26 = v21 + v6 # ADD [1 cycle]\n", + "27: v27 = v26 + v7 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v24 # STORE [4 cycles]\n", + "// TOTAL: 56 cycles\n", + "\n", + "============================================================\n", + "Test 2 (verify fan-out structure):\n", + " Top 5 most-used variables (the 'hubs'):\n", + " v8: used 8 times\n", + " v0: used 5 times\n", + " v9: used 3 times\n", + " v1: used 2 times\n", + " addr_in0: used 1 times\n", + " (In linear Level 3, top variables typically used 1-3 times.)\n", + "\n", + "============================================================\n", + "Test 3 (-O3 baseline on Level 4, seed=42):\n", + "Original (64 cycles):\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v25 = 15 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = 4 # CONST [1 cycle]\n", + "7: v6 = 3 # CONST [1 cycle]\n", + "8: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "10: v9 = v8 // 5 # DIV [5 cycles]\n", + "11: v10 = v8 + v6 # ADD [1 cycle]\n", + "12: v11 = v8 - v2 # SUB [1 cycle]\n", + "13: v12 = v0 // 1 # DIV [5 cycles]\n", + "14: v13 = v9 // 3 # DIV [5 cycles]\n", + "15: v14 = v0 + v1 # ADD [1 cycle]\n", + "16: v23 = 21 # CONST [1 cycle]\n", + "17: v15 = v0 * v9 # MUL [3 cycles]\n", + "18: v16 = v8 - v9 # SUB [1 cycle]\n", + "19: v17 = v12 + v8 # ADD [1 cycle]\n", + "20: v18 = v0 // 2 # DIV [5 cycles]\n", + "21: v24 = 13 # CONST [1 cycle]\n", + "22: v19 = v8 * v5 # MUL [3 cycles]\n", + "23: v20 = v8 - 2 # SUB [1 cycle]\n", + "24: v21 = v0 + 3 # ADD [1 cycle]\n", + "25: v22 = v18 * 2 # MUL [3 cycles]\n", + "26: v26 = v22 + v7 # ADD [1 cycle]\n", + "27: v27 = v26 + v8 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 64 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 64\n", + " optimized_cycles : 19\n", + " reduction_pct : 70.3%\n", + " iterations : 2\n", + "\n", + "============================================================\n", + "Test 4 (-O3 across 20 random Level 4 programs):\n", + " total original cycles : 1102\n", + " total optimized cycles : 345\n", + " avg reduction pct : 68.7%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 5 (held-out test set, Levels 1-4, 20 each = 80 programs):\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig=43.4, avg opt=17.6, agg reduction=59.5%\n", + " Level 4: n=20, avg orig=55.0, avg opt=17.6, agg reduction=68.0%\n", + "\n", + " Overall: n=80, agg reduction=61.4%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_5():\n", + " \"\"\"\n", + " Generate a Level 5 Toy-IR program β€” maximum adversarial density.\n", + "\n", + " Level 5 is the hardest curriculum tier. Programs are long (30-40\n", + " instructions), peppered with dead code, heavy on expensive ops\n", + " (MUL/DIV), and feature multiple LOADs and STOREs. The 70% literal-bait\n", + " rate ensures peephole has many opportunities, and the deep arithmetic\n", + " chains create real differentiation between optimization strategies.\n", + "\n", + " Distinct from Level 4: longer, denser, more peephole bait, more\n", + " expensive operations, multiple observable outputs. The absolute cycle\n", + " savings are huge, magnifying any sub-optimal pass orderings.\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 7-9 CONSTs\n", + " num_consts = random.randint(7, 9)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 20)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: 3 LOADs\n", + " loaded_vars = []\n", + " for i in range(3):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: Designate 3-4 hub variables\n", + " num_hubs = random.randint(3, 4)\n", + " candidate_hubs = const_vars[:5] + loaded_vars\n", + " hub_vars = random.sample(candidate_hubs, min(num_hubs, len(candidate_hubs)))\n", + "\n", + " # Step 4: 15-22 arithmetic ops, heavy on expensive ops + peephole bait\n", + " num_arith = random.randint(15, 22)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " # Bias toward expensive ops (MUL/DIV) for higher cycle savings on optimization\n", + " op = random.choices(\n", + " [\"ADD\", \"SUB\", \"MUL\", \"DIV\"],\n", + " weights=[2, 2, 3, 2], # MUL slightly preferred\n", + " )[0]\n", + "\n", + " # 70% chance of literal src2 β†’ peephole bait\n", + " use_literal_src2 = random.random() < 0.70\n", + "\n", + " # 60% chance src1 is a hub\n", + " if random.random() < 0.60:\n", + " src1 = random.choice(hub_vars)\n", + " else:\n", + " src1 = random.choice(available_vars)\n", + "\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5]) # peephole bait: 1\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3, 4]) # peephole bait: 0, 1, 2\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 5: 5-8 dead CONSTs\n", + " num_dead = random.randint(5, 8)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 50)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 6: Force all 3 LOADed values into the final chain\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load1,\n", + " \"src1\": last_result, \"src2\": loaded_vars[0],\n", + " })\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load2,\n", + " \"src1\": after_load1, \"src2\": loaded_vars[1],\n", + " })\n", + " after_load3 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load3,\n", + " \"src1\": after_load2, \"src2\": loaded_vars[2],\n", + " })\n", + " last_result = after_load3\n", + "\n", + " # Step 7: Multiple STOREs (always 2-3)\n", + " num_stores = random.randint(2, 3)\n", + "\n", + " # Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + " observable_addrs = [0]\n", + "\n", + " # Secondary STOREs use early CONSTs (gives DCE more variety to handle)\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " for i in range(1, num_stores):\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " addr_idx = i # addr1, addr2\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": f\"addr{addr_idx}\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(addr_idx)\n", + "\n", + " # Step 8: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr2\": 2,\n", + " \"addr_in0\": 3,\n", + " \"addr_in1\": 4,\n", + " \"addr_in2\": 5,\n", + " }\n", + " initial_mem = {\n", + " 3: random.randint(1, 50),\n", + " 4: random.randint(1, 50),\n", + " 5: random.randint(1, 50),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 5\n", + "LEVEL_GENERATORS[5] = generate_level_5" + ], + "metadata": { + "id": "vExn6K85Vk7D" + }, + "execution_count": 29, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 5 generator ===\n", + "\n", + "# Test 1: spot-check three seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_5()\n", + " print(f\"\\n=== Level 5, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + "\n", + "# Test 2: baseline on Level 5\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (-O3 baseline on Level 5, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_5()\n", + "result = evaluate_baseline(prog)\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "\n", + "# Test 3: aggregate on Level 5\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 5 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(400, 420):\n", + " random.seed(seed)\n", + " prog = generate_level_5()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 4: full headline plot β€” held-out test set, all 5 levels, 20 each = 100 programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (FULL HELD-OUT TEST SET, Levels 1-5, 100 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3, 4, 5])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "print(f\"\\n === BY LEVEL ===\")\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:6.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:6.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n === OVERALL ===\")\n", + "print(f\" n_programs : {o['n_programs']}\")\n", + "print(f\" total original cycles : {o['total_original']}\")\n", + "print(f\" total optimized cycles : {o['total_optimized']}\")\n", + "print(f\" aggregate reduction pct : {o['aggregate_reduction_pct']:.1%}\")\n", + "print(f\" avg per-program redux : {o['avg_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Igt72oqfWuPV", + "outputId": "d1c12a9e-3c7f-4fd1-e06d-2f414b41b6fd" + }, + "execution_count": 30, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 5, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1], mem[2]\n", + "0: v32 = 36 # CONST [1 cycle]\n", + "1: v0 = 4 # CONST [1 cycle]\n", + "2: v1 = 1 # CONST [1 cycle]\n", + "3: v2 = 9 # CONST [1 cycle]\n", + "4: v3 = 8 # CONST [1 cycle]\n", + "5: v4 = 8 # CONST [1 cycle]\n", + "6: v5 = 5 # CONST [1 cycle]\n", + "7: v33 = 44 # CONST [1 cycle]\n", + "8: v6 = 4 # CONST [1 cycle]\n", + "9: v7 = 18 # CONST [1 cycle]\n", + "10: v8 = 3 # CONST [1 cycle]\n", + "11: v9 = mem[addr_in0] # LOAD [4 cycles]\n", + "12: v10 = mem[addr_in1] # LOAD [4 cycles]\n", + "13: v11 = mem[addr_in2] # LOAD [4 cycles]\n", + "14: v12 = v1 * 1 # MUL [3 cycles]\n", + "15: v13 = v12 * 1 # MUL [3 cycles]\n", + "16: v14 = v10 * 0 # MUL [3 cycles]\n", + "17: v35 = 49 # CONST [1 cycle]\n", + "18: v37 = 8 # CONST [1 cycle]\n", + "19: v15 = v9 + 2 # ADD [1 cycle]\n", + "20: v31 = 30 # CONST [1 cycle]\n", + "21: v16 = v1 // 5 # DIV [5 cycles]\n", + "22: v36 = 50 # CONST [1 cycle]\n", + "23: v17 = v11 - 4 # SUB [1 cycle]\n", + "24: v18 = v9 + 0 # ADD [1 cycle]\n", + "25: v30 = 25 # CONST [1 cycle]\n", + "26: v19 = v1 // 2 # DIV [5 cycles]\n", + "27: v20 = v0 - 4 # SUB [1 cycle]\n", + "28: v21 = v1 * 3 # MUL [3 cycles]\n", + "29: v22 = v7 - v21 # SUB [1 cycle]\n", + "30: v23 = v0 - v10 # SUB [1 cycle]\n", + "31: v24 = v18 - 2 # SUB [1 cycle]\n", + "32: v25 = v20 + 3 # ADD [1 cycle]\n", + "33: v26 = v17 + 2 # ADD [1 cycle]\n", + "34: v27 = v10 * 1 # MUL [3 cycles]\n", + "35: v28 = v0 // 1 # DIV [5 cycles]\n", + "36: v29 = v19 + 0 # ADD [1 cycle]\n", + "37: v34 = 44 # CONST [1 cycle]\n", + "38: v38 = v29 + v9 # ADD [1 cycle]\n", + "39: v39 = v38 + v10 # ADD [1 cycle]\n", + "40: v40 = v39 + v11 # ADD [1 cycle]\n", + "41: mem[addr0] = v40 # STORE [4 cycles]\n", + "42: mem[addr1] = v1 # STORE [4 cycles]\n", + "43: mem[addr2] = v33 # STORE [4 cycles]\n", + "// TOTAL: 84 cycles\n", + "\n", + "=== Level 5, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v34 = 29 # CONST [1 cycle]\n", + "1: v0 = 19 # CONST [1 cycle]\n", + "2: v33 = 6 # CONST [1 cycle]\n", + "3: v1 = 3 # CONST [1 cycle]\n", + "4: v2 = 9 # CONST [1 cycle]\n", + "5: v3 = 4 # CONST [1 cycle]\n", + "6: v32 = 3 # CONST [1 cycle]\n", + "7: v4 = 16 # CONST [1 cycle]\n", + "8: v5 = 15 # CONST [1 cycle]\n", + "9: v6 = 16 # CONST [1 cycle]\n", + "10: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "11: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "12: v9 = mem[addr_in2] # LOAD [4 cycles]\n", + "13: v10 = v8 - v4 # SUB [1 cycle]\n", + "14: v11 = v5 * 0 # MUL [3 cycles]\n", + "15: v12 = v8 + 1 # ADD [1 cycle]\n", + "16: v13 = v8 // 5 # DIV [5 cycles]\n", + "17: v36 = 16 # CONST [1 cycle]\n", + "18: v35 = 49 # CONST [1 cycle]\n", + "19: v14 = v8 - 2 # SUB [1 cycle]\n", + "20: v15 = v14 // 1 # DIV [5 cycles]\n", + "21: v31 = 6 # CONST [1 cycle]\n", + "22: v16 = v3 + v10 # ADD [1 cycle]\n", + "23: v17 = v8 // 2 # DIV [5 cycles]\n", + "24: v18 = v16 - 3 # SUB [1 cycle]\n", + "25: v19 = v8 * 3 # MUL [3 cycles]\n", + "26: v20 = v11 * 0 # MUL [3 cycles]\n", + "27: v21 = v16 - 3 # SUB [1 cycle]\n", + "28: v22 = v9 - v19 # SUB [1 cycle]\n", + "29: v23 = v0 * 0 # MUL [3 cycles]\n", + "30: v24 = v7 * 3 # MUL [3 cycles]\n", + "31: v25 = v8 * v8 # MUL [3 cycles]\n", + "32: v26 = v12 * 4 # MUL [3 cycles]\n", + "33: v27 = v8 // 1 # DIV [5 cycles]\n", + "34: v28 = v8 * 3 # MUL [3 cycles]\n", + "35: v29 = v9 // 4 # DIV [5 cycles]\n", + "36: v30 = v17 * v18 # MUL [3 cycles]\n", + "37: v37 = v30 + v7 # ADD [1 cycle]\n", + "38: v38 = v37 + v8 # ADD [1 cycle]\n", + "39: v39 = v38 + v9 # ADD [1 cycle]\n", + "40: mem[addr0] = v39 # STORE [4 cycles]\n", + "41: mem[addr1] = v32 # STORE [4 cycles]\n", + "// TOTAL: 95 cycles\n", + "\n", + "=== Level 5, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1], mem[2]\n", + "0: v0 = 5 # CONST [1 cycle]\n", + "1: v1 = 13 # CONST [1 cycle]\n", + "2: v2 = 2 # CONST [1 cycle]\n", + "3: v3 = 3 # CONST [1 cycle]\n", + "4: v4 = 18 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 12 # CONST [1 cycle]\n", + "7: v7 = 19 # CONST [1 cycle]\n", + "8: v8 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v34 = 1 # CONST [1 cycle]\n", + "10: v9 = mem[addr_in1] # LOAD [4 cycles]\n", + "11: v10 = mem[addr_in2] # LOAD [4 cycles]\n", + "12: v32 = 32 # CONST [1 cycle]\n", + "13: v11 = v3 - 4 # SUB [1 cycle]\n", + "14: v12 = v0 + 4 # ADD [1 cycle]\n", + "15: v13 = v9 * 1 # MUL [3 cycles]\n", + "16: v14 = v0 - 4 # SUB [1 cycle]\n", + "17: v15 = v9 // 2 # DIV [5 cycles]\n", + "18: v16 = v3 - 4 # SUB [1 cycle]\n", + "19: v33 = 17 # CONST [1 cycle]\n", + "20: v17 = v0 + 3 # ADD [1 cycle]\n", + "21: v18 = v3 * 1 # MUL [3 cycles]\n", + "22: v19 = v0 + 3 # ADD [1 cycle]\n", + "23: v20 = v0 - v5 # SUB [1 cycle]\n", + "24: v21 = v3 * 0 # MUL [3 cycles]\n", + "25: v22 = v10 * 2 # MUL [3 cycles]\n", + "26: v23 = v0 * 0 # MUL [3 cycles]\n", + "27: v24 = v9 // 1 # DIV [5 cycles]\n", + "28: v25 = v9 * 3 # MUL [3 cycles]\n", + "29: v26 = v0 - 3 # SUB [1 cycle]\n", + "30: v27 = v3 - 2 # SUB [1 cycle]\n", + "31: v28 = v0 + 0 # ADD [1 cycle]\n", + "32: v29 = v3 + 3 # ADD [1 cycle]\n", + "33: v30 = v0 // 4 # DIV [5 cycles]\n", + "34: v35 = 27 # CONST [1 cycle]\n", + "35: v31 = v3 // 2 # DIV [5 cycles]\n", + "36: v36 = 24 # CONST [1 cycle]\n", + "37: v37 = v31 + v8 # ADD [1 cycle]\n", + "38: v38 = v37 + v9 # ADD [1 cycle]\n", + "39: v39 = v38 + v10 # ADD [1 cycle]\n", + "40: mem[addr0] = v39 # STORE [4 cycles]\n", + "41: mem[addr1] = v2 # STORE [4 cycles]\n", + "42: mem[addr2] = v0 # STORE [4 cycles]\n", + "// TOTAL: 89 cycles\n", + "\n", + "============================================================\n", + "Test 2 (-O3 baseline on Level 5, seed=42):\n", + " original_cycles : 84\n", + " optimized_cycles : 29\n", + " reduction_pct : 65.5%\n", + " iterations : 2\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 5 programs):\n", + " total original cycles : 1659\n", + " total optimized cycles : 534\n", + " avg reduction pct : 67.8%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 4 (FULL HELD-OUT TEST SET, Levels 1-5, 100 programs):\n", + "\n", + " === BY LEVEL ===\n", + " Level 1: n=20, avg orig= 9.9, avg opt= 5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig= 25.2, avg opt= 11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig= 43.4, avg opt= 17.6, agg reduction=59.5%\n", + " Level 4: n=20, avg orig= 55.0, avg opt= 17.6, agg reduction=68.0%\n", + " Level 5: n=20, avg orig= 84.0, avg opt= 26.1, agg reduction=68.9%\n", + "\n", + " === OVERALL ===\n", + " n_programs : 100\n", + " total original cycles : 4350\n", + " total optimized cycles : 1555\n", + " aggregate reduction pct : 64.3%\n", + " avg per-program redux : 59.3%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [], + "metadata": { + "id": "iCDDLZScWvlr" + }, + "execution_count": 30, + "outputs": [] + } + ] +} \ No newline at end of file diff --git a/space/space/space/space/space/space/space/space/space/app.py b/space/space/space/space/space/space/space/space/space/app.py index 061e42b732d69ec39765c751e844799b5df53f17..75ac67126c7d94654b416aaa1d52b9610ab7b14a 100644 --- a/space/space/space/space/space/space/space/space/space/app.py +++ b/space/space/space/space/space/space/space/space/space/app.py @@ -23,6 +23,7 @@ from runtime_core import ( MockEngine, SAMPLE_PROGRAM, ) +from train import run_toy_training DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2) @@ -113,8 +114,9 @@ def build_demo() -> gr.Blocks: # Compiler optimization (Toy-IR) β€” interactive demo This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing, - plus the **CompilerOptimizationEnv** smoke loop from your notebooks. Full GRPO / Unsloth - training belongs on Colab or a GPU Space. + plus the **CompilerOptimizationEnv** loop. A **toy REINFORCE** tab trains a tiny + stateless policy over the mock passes (see `train.py`). Full **GRPO + LLM + Unsloth** + still belongs on Colab or a GPU machine. """ ).strip() ) @@ -154,6 +156,28 @@ def build_demo() -> gr.Blocks: out_ep = gr.Textbox(label="Log", lines=20) gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep]) + with gr.Tab("Toy training (REINFORCE)"): + gr.Markdown( + textwrap.dedent( + """ + Trains a **stateless** categorical policy over the three mock passes + (`constant_folding`, `dead_code_elimination`, `loop_unrolling`) using + **REINFORCE** in pure Python. Same logic as: `python train.py --episodes 50`. + + This is a CPU smoke run, not an LLM. For real GRPO, use your project notebooks + on a GPU. + """ + ).strip() + ) + tr_ep = gr.Slider(5, 200, value=50, step=1, label="episodes") + tr_ms = gr.Slider(2, 20, value=8, step=1, label="max_steps per episode") + tr_seed = gr.Number(value=0, label="random seed", precision=0) + tr_lr = gr.Slider(0.01, 0.5, value=0.15, step=0.01, label="learning rate") + out_tr = gr.Textbox(label="Training log", lines=18) + gr.Button("Run training", variant="primary").click( + run_toy_training, [tr_ep, tr_ms, tr_seed, tr_lr], [out_tr] + ) + gr.Markdown( "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, " "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`." diff --git a/space/space/space/space/space/space/space/space/space/space/space/runtime_core.py b/space/space/space/space/space/space/space/space/space/space/space/runtime_core.py new file mode 100644 index 0000000000000000000000000000000000000000..75c7b62e4d35ae660c4f4a3fec17e17beb79a965 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/runtime_core.py @@ -0,0 +1,316 @@ +""" +Shared runtime for the Gradio Space: mock compiler env + Deliverable 2 formatting. +Sourced from `compiler_optimization_grpo.ipynb` and +`role2_deliverable3_training_loop (2) (1) (1).ipynb`. +""" + +from __future__ import annotations + +import copy +import json +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +# --- Deliverable 2 (LLM-facing pseudo-asm + pass-array parsing) ----------------- + + +class Deliverable2_Formatter: + @staticmethod + def translate_state(raw_json: list) -> str: + """Translate raw JSON IR into compact pseudo-assembly.""" + if not isinstance(raw_json, list) or not raw_json: + return "; (empty program β€” 0 instructions)" + + pseudo_assembly: list[str] = [] + for i, instruction in enumerate(raw_json): + if not isinstance(instruction, dict): + pseudo_assembly.append(f"{i}. NOP") + continue + + op = str(instruction.get("op", "UNKNOWN")).upper() + args = ", ".join(str(arg) for arg in instruction.get("args", [])) + dest = instruction.get("dest", "") + if dest: + line = f"{i}. {dest} = {op} {args}".rstrip() + else: + line = f"{i}. {op} {args}".rstrip() + pseudo_assembly.append(line) + + return "\n".join(pseudo_assembly) + + @staticmethod + def extract_action_array(llm_output: str) -> list: + """Best-effort extraction of JSON pass arrays from noisy LLM output.""" + text = (llm_output or "").strip() + if not text: + raise ValueError("Invalid JSON format") + + try: + parsed = json.loads(text) + if isinstance(parsed, list): + return parsed + except json.JSONDecodeError: + pass + + cleaned = re.sub(r"```(?:json)?", "", text, flags=re.IGNORECASE).replace("```", "").strip() + if cleaned != text: + try: + parsed = json.loads(cleaned) + if isinstance(parsed, list): + return parsed + except json.JSONDecodeError: + pass + + match = re.search(r"\[.*?\]", text, re.DOTALL) + if match: + candidate = match.group(0) + try: + parsed = json.loads(candidate) + if isinstance(parsed, list): + return parsed + except json.JSONDecodeError: + try: + parsed = json.loads(candidate.replace("'", '"')) + if isinstance(parsed, list): + return parsed + except json.JSONDecodeError: + pass + + raise ValueError("Invalid JSON format") + + +# --- OpenEnv-style compiler environment (mock engine) -------------------------- + + +class MCPEnvironment: + """Minimal stub. In production: `from openenv import MCPEnvironment`.""" + + def reset(self, *args, **kwargs): + raise NotImplementedError + + def step(self, *args, **kwargs): + raise NotImplementedError + + def state(self): + raise NotImplementedError + + +@dataclass +class StepResult: + observation: str + reward: float + done: bool + info: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class EpisodeStats: + steps_taken: int = 0 + total_reward: float = 0.0 + passes_applied: List[str] = field(default_factory=list) + invalid_actions: int = 0 + no_ops: int = 0 + baseline_cycles: int = 0 + final_cycles: int = 0 + + @property + def total_improvement_pct(self) -> float: + if self.baseline_cycles == 0: + return 0.0 + return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0 + + +class CompilerOptimizationEnv(MCPEnvironment): + TIME_TAX: float = 1.0 + NO_OP_PENALTY: float = -2.0 + INVALID_ACTION_PENALTY: float = -5.0 + MAX_INVALID_ACTIONS: int = 3 + TERMINAL_BONUS_SCALE: float = 0.5 + + def __init__( + self, + role1_engine, + role3_passes: Dict[str, Any], + max_steps: int = 10, + curriculum_level: int = 1, + ): + self.engine = role1_engine + self.passes = role3_passes + self.max_steps = max_steps + self.curriculum_level = curriculum_level + self._valid_actions = frozenset(self.passes.keys()) + + self._stats: Optional[EpisodeStats] = None + self.original_program = None + self.current_program = None + self.previous_cycles = 0 + self._consecutive_invalid = 0 + + def reset(self, new_program_json: List[Dict]) -> str: + self.original_program = copy.deepcopy(new_program_json) + self.current_program = copy.deepcopy(new_program_json) + self.previous_cycles = self._safe_count_cycles(self.current_program) + self._consecutive_invalid = 0 + self._stats = EpisodeStats( + baseline_cycles=self.previous_cycles, + final_cycles=self.previous_cycles, + ) + return self.state() + + def state(self) -> str: + assert self.current_program is not None + return self._program_to_pseudoasm(self.current_program) + + def step(self, action_string: str) -> StepResult: + assert self._stats is not None, "Call reset() before step()." + self._stats.steps_taken += 1 + + if action_string not in self._valid_actions: + return self._handle_invalid_action(action_string) + + candidate_program = self.passes[action_string](copy.deepcopy(self.current_program)) + + is_valid = self.engine.verify_equivalence(self.original_program, candidate_program) + if not is_valid: + return self._handle_semantic_violation() + + new_cycles = self._safe_count_cycles(candidate_program) + reward, info = self._compute_reward(action_string, new_cycles) + + self.current_program = candidate_program + self.previous_cycles = new_cycles + self._stats.final_cycles = new_cycles + self._stats.total_reward += reward + self._stats.passes_applied.append(action_string) + self._consecutive_invalid = 0 + + done = self._stats.steps_taken >= self.max_steps + if done: + terminal_bonus = self._terminal_bonus() + reward += terminal_bonus + info["terminal_bonus"] = terminal_bonus + info["reason"] = "max_steps_reached" + info["episode_stats"] = self._episode_summary() + + return StepResult(self.state(), reward, done, info) + + def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, Dict]: + info: Dict[str, Any] = {"action": action} + if self.previous_cycles == 0: + return -self.TIME_TAX, {**info, "note": "zero_baseline"} + + old_cycles = self.previous_cycles + delta_pct = ((old_cycles - new_cycles) / old_cycles) * 100.0 + + if new_cycles == old_cycles: + reward = self.NO_OP_PENALTY + if self._stats is not None: + self._stats.no_ops += 1 + info["no_op"] = True + else: + reward = delta_pct - self.TIME_TAX + info["delta_pct"] = round(delta_pct, 3) + + info["prev_cycles"] = old_cycles + info["new_cycles"] = new_cycles + return reward, info + + def _terminal_bonus(self) -> float: + if self._stats is None: + return 0.0 + return max(0.0, self._stats.total_improvement_pct * self.TERMINAL_BONUS_SCALE) + + def _compute_crash_penalty(self) -> float: + return -2.0 * (100.0 * self.max_steps) + + def _handle_invalid_action(self, action: str) -> StepResult: + self._consecutive_invalid += 1 + if self._stats is not None: + self._stats.invalid_actions += 1 + done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS + info = { + "error": f"Unknown action: '{action}'", + "valid_actions": sorted(self._valid_actions), + "consecutive_invalid": self._consecutive_invalid, + } + if done: + info["reason"] = "too_many_invalid_actions" + info["episode_stats"] = self._episode_summary() + return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info) + + def _handle_semantic_violation(self) -> StepResult: + return StepResult( + self.state(), + self._compute_crash_penalty(), + True, + { + "error": "Semantic equivalence check FAILED.", + "reason": "semantic_violation", + "episode_stats": self._episode_summary(), + }, + ) + + @staticmethod + def _program_to_pseudoasm(program: List[Dict]) -> str: + if not program: + return "; (empty program)" + lines = [] + for i, instr in enumerate(program): + op = instr.get("op", "NOP") + args = instr.get("args", []) + dest = instr.get("dest") + typ = instr.get("type", "") + arg_str = ", ".join(str(a) for a in args) + type_hint = f":{typ}" if typ else "" + if dest: + lines.append(f" {i:>3}: {dest}{type_hint} = {op} {arg_str}") + else: + lines.append(f" {i:>3}: {op} {arg_str}") + return "\n".join(lines) + + def _safe_count_cycles(self, program: List[Dict]) -> int: + return max(0, int(self.engine.execute_and_count_cycles(program))) + + def _episode_summary(self) -> Dict: + s = self._stats + if s is None: + return {} + return { + "steps": s.steps_taken, + "total_reward": round(s.total_reward, 3), + "passes_applied": s.passes_applied, + "invalid_actions": s.invalid_actions, + "no_ops": s.no_ops, + "baseline_cycles": s.baseline_cycles, + "final_cycles": s.final_cycles, + "total_improvement_pct": round(s.total_improvement_pct, 3), + } + + def available_actions(self) -> List[str]: + return sorted(self._valid_actions) + + +class MockEngine: + """Stub engine: cycles = instruction count, all programs semantically valid.""" + + def execute_and_count_cycles(self, program): + return len(program) + + def verify_equivalence(self, original, candidate): + return True + + +MOCK_PASSES = { + "constant_folding": lambda p: p[:-1] if len(p) > 1 else p, + "dead_code_elimination": lambda p: p[:-1] if len(p) > 2 else p, + "loop_unrolling": lambda p: p, +} + +SAMPLE_PROGRAM = [ + {"op": "const", "dest": "x", "args": ["5"], "type": "int"}, + {"op": "const", "dest": "y", "args": ["3"], "type": "int"}, + {"op": "add", "dest": "z", "args": ["x", "y"], "type": "int"}, + {"op": "mul", "dest": "w", "args": ["z", "x"], "type": "int"}, + {"op": "ret", "args": ["w"]}, +] diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/Dockerfile b/space/space/space/space/space/space/space/space/space/space/space/space/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..445eb4cfc5eec77c8dbf2ff58f3803f047e6e6ee --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/Dockerfile @@ -0,0 +1,19 @@ +# Optional: use Hugging Face Space SDK = docker (set `sdk: docker` in README.md). +# Default README uses Gradio SDK and does not require this image. + +FROM python:3.11-slim + +WORKDIR /app + +ENV PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + GRADIO_SERVER_NAME=0.0.0.0 + +COPY requirements.txt . +RUN pip install --upgrade pip && pip install -r requirements.txt + +COPY . . + +EXPOSE 7860 + +CMD ["python", "app.py"] diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/README.md b/space/space/space/space/space/space/space/space/space/space/space/space/README.md index 55dcc433c9908155a8f5fb93a51169b100921b85..81bde60282565f4f9bb8b67839a17bd997c684eb 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/README.md +++ b/space/space/space/space/space/space/space/space/space/space/space/space/README.md @@ -1,9 +1,54 @@ --- -title: Compiler Brain OpenEnv -emoji: 🧠 -colorFrom: indigo -colorTo: purple -sdk: docker +title: Compiler Tetris β€” Toy-IR RL Demo +emoji: 🧩 +colorFrom: gray +colorTo: indigo +sdk: gradio +sdk_version: 5.12.0 app_file: app.py pinned: false +license: mit --- + +# Compiler Tetris (Toy-IR) β€” Hugging Face Space + +Interactive **CPU** demo for the Meta / OpenEnv-style **compiler phase-ordering** project: + +- **Deliverable 2 β€” state translation:** JSON Toy-IR β†’ compact pseudo-assembly for the LLM. +- **Deliverable 2 β€” format enforcement:** resilient extraction of a JSON **pass array** from noisy model text. +- **CompilerOptimizationEnv** (mock Role 1 engine + mock passes): step through passes and inspect rewards. + +Full **GRPO / Unsloth** training is not run here (heavy GPU + long installs). Use your notebooks on Colab or a GPU Space for training. + +## Source notebooks (parent repo) + +- `compiler_optimization_grpo.ipynb` +- `role2_deliverable3_training_loop (2) (1).ipynb` +- `compiler_tetris (1).ipynb` +- `metahack1 (1).ipynb` + +## Deploy this folder as a new Space + +1. Create a new Space on Hugging Face (SDK: **Gradio**). +2. Upload the contents of this `hf_space/` directory to the Space repository root (`app.py`, `requirements.txt`, `README.md`). +3. Optional: add `Dockerfile` only if you switch the Space to **Docker** (see below). + +## Optional: Docker SDK instead of Gradio SDK + +If you want the Space to build from `Dockerfile`, change the YAML header to: + +```yaml +sdk: docker +``` + +and remove Gradio-specific keys (`sdk_version`, `app_file`). The container runs `python app.py`, which listens on the `PORT` environment variable provided by Spaces. + +## Local run + +```bash +cd hf_space +pip install -r requirements.txt +python app.py +``` + +Then open `http://127.0.0.1:7860`. diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/app.py b/space/space/space/space/space/space/space/space/space/space/space/space/app.py index e8c748e0423b20310f1295fc7c849bd997dd807e..061e42b732d69ec39765c751e844799b5df53f17 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/app.py +++ b/space/space/space/space/space/space/space/space/space/space/space/space/app.py @@ -1,182 +1,167 @@ +""" +Hugging Face Spaces entrypoint β€” Gradio UI for Toy-IR compiler RL demo. +""" + +from __future__ import annotations + import json -import logging -import traceback -from typing import Any, Dict, List - -# OpenEnv SDK import -from openenv import MCPEnvironment - -# ============================================================================== -# ROLE 1 & 3 IMPORTS -# TODO: Import the actual execution engine and generator from your teammates -# ============================================================================== -# from engine import execute_tac, verify_equivalence -# from curriculum import generate_level_code - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("CompilerEnvServer") - -class CompilerEnv(MCPEnvironment): - """ - The OpenEnv Server Wrapper for the Toy-IR Compiler Pass Optimizer. - Acts as the referee between the LLM client and Role 1's Execution Engine. - """ - - def __init__(self): - super().__init__() - # State variables - self.raw_json_code = None - self.current_state_string = "" - self.initial_cycles = 0 - self.current_step = 0 - self.max_steps = 10 # Set a max step limit per episode - - # Cycle Weights (The Physics defined by Role 2) - self.cycle_weights = { - "ADD": 1, - "SUB": 1, - "MUL": 4, - "DIV": 10, - "MEM_LOAD": 20, - "STORE": 20 - } - - def reset(self) -> str: - """ - Grabs unoptimized code, calculates baseline cycles, and translates the - state for the LLM. - """ - self.current_step = 0 - - # 1. Grab new unoptimized code (Role 3 integration) - # TODO: Replace with real generator: self.raw_json_code = generate_level_code() - self.raw_json_code = self._mock_generator() - - # 2. Get baseline cycles (Role 1 integration) - # TODO: Replace with real engine: self.initial_cycles, _ = execute_tac(self.raw_json_code, []) - self.initial_cycles = 100 - - # 3. Translate to Pseudo-Assembly to prevent Attention Dilution - self.current_state_string = self._translate_state(self.raw_json_code) - - logger.info(f"Environment Reset. Baseline Cycles: {self.initial_cycles}") - return self.state() - - def step(self, action: str) -> Dict[str, Any]: - """ - Executes the LLM's chosen optimization pass, verifies math equivalence, - and calculates the reward. - """ - self.current_step += 1 - - # 1. Parse LLM Action (Regex/JSON robustness) - try: - # Assuming the LLM outputs a single pass name or a list of passes - action_data = json.loads(action) - if isinstance(action_data, str): - action_array = [action_data] - else: - action_array = action_data - format_bonus = 0.1 - except json.JSONDecodeError: - # Format Trap - return self._build_step_response( - reward=-2.5, - done=True, - error="Invalid JSON. You must output a valid JSON array of strings." - ) - - # 2. Execute Code & Verify Equivalence (Role 1 Integration) - # TODO: new_cycles, optimized_code = execute_tac(self.raw_json_code, action_array) - # TODO: is_valid = verify_equivalence(self.raw_json_code, optimized_code) - new_cycles = 80 # Mock Data - is_valid = True # Mock Data - - # 3. Calculate Reward Physics - if not is_valid: - # Correctness Penalty + Micro-Variance for GRPO - penalty = -2.0 - (len(action_array) * 0.01) - return self._build_step_response( - reward=penalty, - done=True, - error=f"Code equivalence broken by passes: {action_array}" - ) - - # Calculate Improvement Ratio + Time Tax (-1.0) - cycle_improvement_ratio = (self.initial_cycles - new_cycles) / self.initial_cycles - time_tax = -0.05 * self.current_step # Small tax to prevent pass spamming - reward = cycle_improvement_ratio + format_bonus + time_tax - - # Update state if sequential, or finish if one-shot - # NOTE: For hackathon speed, we treat this as a One-Shot episode - done = True - - return self._build_step_response( - reward=reward, - done=done, - info={"status": "success", "optimized_cycles": new_cycles} +import os +import sys +import textwrap +from pathlib import Path + +import gradio as gr + +_ROOT = Path(__file__).resolve().parent +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from runtime_core import ( + CompilerOptimizationEnv, + Deliverable2_Formatter, + MOCK_PASSES, + MockEngine, + SAMPLE_PROGRAM, +) + + +DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2) + + +def translate_ir(ir_json: str) -> tuple[str, str]: + try: + data = json.loads(ir_json.strip() or "[]") + except json.JSONDecodeError as e: + return "", f"Invalid JSON: {e}" + if not isinstance(data, list): + return "", "JSON root must be a list of instruction objects." + d2 = Deliverable2_Formatter.translate_state(data) + env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10) + env.reset(data) + internal = env.state() + return d2, internal + + +def parse_llm_output(llm_text: str) -> str: + try: + arr = Deliverable2_Formatter.extract_action_array(llm_text) + return json.dumps(arr, indent=2) + except ValueError as e: + return str(e) + + +def parse_action_line(line: str) -> list[str]: + line = line.strip() + if not line: + return [] + try: + got = Deliverable2_Formatter.extract_action_array(line) + return [str(x).strip() for x in got] + except ValueError: + parts = [p.strip().strip("\"'") for p in line.split(",") if p.strip()] + return [p.lower() for p in parts] + + +def run_episode(ir_json: str, actions_multiline: str, max_steps: int) -> str: + try: + program = json.loads(ir_json.strip() or "[]") + except json.JSONDecodeError as e: + return f"Invalid program JSON: {e}" + if not isinstance(program, list): + return "Program must be a JSON list." + + lines = [ln for ln in actions_multiline.splitlines() if ln.strip()] + actions: list[str] = [] + for ln in lines: + actions.extend(parse_action_line(ln)) + if not actions: + return "No actions parsed. Enter JSON arrays or comma-separated pass names." + + engine = MockEngine() + env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=int(max_steps)) + obs0 = env.reset(program) + log = [ + f"Baseline cycles: {env.previous_cycles}", + f"Initial observation (env):\n{obs0}", + "", + f"Deliverable 2 pseudo-assembly:\n{Deliverable2_Formatter.translate_state(program)}", + "", + "--- steps ---", + ] + for i, act in enumerate(actions, start=1): + r = env.step(act) + log.append( + f"{i}. {act!r} β†’ reward={r.reward:+.3f} done={r.done} cycles={env.previous_cycles}" + ) + if r.info: + slim = {k: v for k, v in r.info.items() if k in ("delta_pct", "error", "no_op", "reason", "terminal_bonus")} + if slim: + log.append(f" info: {slim}") + if r.done: + break + log.append("") + log.append("Episode summary:") + log.append(json.dumps(env._episode_summary(), indent=2)) + return "\n".join(log) + + +def build_demo() -> gr.Blocks: + with gr.Blocks(title="Compiler Tetris β€” Toy-IR RL Demo") as demo: + gr.Markdown( + textwrap.dedent( + """ + # Compiler optimization (Toy-IR) β€” interactive demo + + This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing, + plus the **CompilerOptimizationEnv** smoke loop from your notebooks. Full GRPO / Unsloth + training belongs on Colab or a GPU Space. + """ + ).strip() + ) + + with gr.Tabs(): + with gr.Tab("Translate IR (D2)"): + gr.Markdown("Paste Toy-IR as a JSON **list** of instruction dicts (`op`, `args`, `dest`, …).") + ir_in = gr.Textbox(label="IR JSON", value=DEFAULT_IR, lines=12, max_lines=24) + btn_t = gr.Button("Translate", variant="primary") + out_d2 = gr.Textbox(label="Deliverable2 pseudo-assembly (LLM-facing)", lines=10) + out_env = gr.Textbox(label="Env internal pseudo-asm (with type hints)", lines=10) + btn_t.click(translate_ir, [ir_in], [out_d2, out_env]) + + with gr.Tab("Parse LLM output (D2)"): + gr.Markdown( + "Paste model output. A JSON array is preferred; bracket extraction handles light noise." + ) + llm_in = gr.Textbox( + label="LLM output", + value='Here is the plan: ["constant_folding", "dead_code_elimination"]', + lines=4, + ) + out_parse = gr.Textbox(label="Parsed array or error", lines=6) + gr.Button("Parse", variant="primary").click(parse_llm_output, [llm_in], [out_parse]) + + with gr.Tab("Env episode (mock engine)"): + gr.Markdown( + "One JSON program + actions: each line can be a JSON array or comma-separated names." + ) + ir_ep = gr.Textbox(label="Program JSON", value=DEFAULT_IR, lines=10) + acts = gr.Textbox( + label="Actions (one JSON array or comma-list per line)", + value='["constant_folding", "dead_code_elimination"]\nloop_unrolling', + lines=5, + ) + ms = gr.Slider(1, 20, value=10, step=1, label="max_steps") + out_ep = gr.Textbox(label="Log", lines=20) + gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep]) + + gr.Markdown( + "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, " + "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`." ) - def state(self) -> str: - """ - Returns the current observation to the LLM. - """ - return f"Current Step: {self.current_step}/{self.max_steps}\n\n{self.current_state_string}" - - def _translate_state(self, raw_json: List[Dict]) -> str: - """ - Translates raw AST JSON into clean pseudo-assembly. - Strips all UUIDs and AST metadata. - """ - pseudo_assembly = [] - instruction_count = 1 - - for inst in raw_json: - op = inst.get("op", "UNKNOWN") - src1 = inst.get("src1", "") - src2 = inst.get("src2", "") - dest = inst.get("dest", "") - - # Format arguments cleanly - args = f"{src1}" if src2 is None else f"{src1}, {src2}" - - if dest: - line = f"{instruction_count}. {dest} = {op} {args}" - else: - line = f"{instruction_count}. {op} {args}" - - pseudo_assembly.append(line) - instruction_count += 1 - - return "\n".join(pseudo_assembly) - - def _build_step_response(self, reward: float, done: bool, error: str = None, info: dict = None) -> Dict[str, Any]: - """Helper to format the standard OpenEnv step return dictionary.""" - response = { - "reward": reward, - "done": done, - "state": self.state() - } - if error: - response["error"] = error - if info: - response["info"] = info - return response - - def _mock_generator(self): - """Mock data so the server runs before Role 1 integrates their engine.""" - return [ - {"op": "CONST", "dest": "a", "src1": 2, "src2": None}, - {"op": "CONST", "dest": "b", "src1": 3, "src2": None}, - {"op": "ADD", "dest": "c", "src1": "a", "src2": "b"} - ] + return demo + if __name__ == "__main__": - logger.info("Initializing CompilerEnv Server...") - try: - env = CompilerEnv() - # openenv.run() or start() depending on the specific MCP wrapper version - env.start() - except Exception as e: - logger.error(f"Failed to start environment server: {e}") - logger.error(traceback.format_exc()) + port = int(os.environ.get("PORT", "7860")) + build_demo().launch(server_name="0.0.0.0", server_port=port) diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt b/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt index 4ca6c261121de442a517f2fc7ae5ce02fb6b7ac6..2bd52bdfe7e82c665e1a3ce6325f8259bfa1c566 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt +++ b/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt @@ -1,24 +1,2 @@ -# OpenEnv SDK (Mandatory for the hackathon judging environment) -openenv - -# Core RL Environment -gymnasium - -# Hugging Face Training Stack (Versions matched to your D3 notebook) -trl==0.23.1 -transformers==4.57.1 -peft -accelerate -bitsandbytes - -# Unsloth for efficient 4-bit QLoRA training on T4 GPUs -# Note: Unsloth often prefers being installed via their specific pip wheel or git, -# but this is standard for a requirements file. -unsloth - -# PyTorch (Will default to standard compatible version if no index is specified) -torch - -# Logging & Visualizations (For Day 2 judging criteria) -wandb -matplotlib +# Hugging Face Spaces β€” Gradio SDK (CPU demo) +gradio>=4.44.0,<6 diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_optimization_grpo.ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_optimization_grpo.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..16c0a6e9dbd3f1af944c4c66ad953d6f173daaca --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_optimization_grpo.ipynb @@ -0,0 +1,959 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + }, + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "markdown", + "id": "title-cell", + "metadata": {}, + "source": [ + "# πŸ”§ Compiler Optimization RL Environment\n", + "### OpenEnv Hackathon 2026 β€” Theme #2: Long-Horizon Planning\n", + "\n", + "**What this notebook does:**\n", + "1. Defines a fully OpenEnv-compliant `CompilerOptimizationEnv`\n", + "2. Loads `Qwen2.5-3B-Instruct` via Unsloth (4-bit QLoRA)\n", + "3. Trains with GRPO (TRL) β€” LLM learns to pick compiler passes that reduce CPU cycles\n", + "4. Runs a smoke-test with a mock engine so you can verify reward logic without real hardware\n", + "5. Plots reward curves\n", + "\n", + "---\n", + "**Stack:** `unsloth` Β· `trl` Β· `openenv` Β· `wandb` Β· `matplotlib`\n", + "\n", + "> **Runtime:** Google Colab T4 GPU recommended. For the smoke-test only, CPU is fine." + ] + }, + { + "cell_type": "markdown", + "id": "install-header", + "metadata": {}, + "source": [ + "## πŸ“¦ Cell 1 β€” Install Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "install-cell", + "metadata": {}, + "outputs": [], + "source": [ + "# Install all required packages\n", + "# Unsloth must be installed before trl to get the right CUDA kernels\n", + "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n", + "!pip install trl datasets transformers accelerate peft bitsandbytes --quiet\n", + "!pip install wandb matplotlib --quiet\n", + "\n", + "# Optional: install openenv for production use\n", + "# !pip install openenv --quiet\n", + "\n", + "print(\"βœ… All packages installed\")" + ] + }, + { + "cell_type": "markdown", + "id": "env-header", + "metadata": {}, + "source": [ + "## 🌍 Cell 2 β€” CompilerOptimizationEnv (OpenEnv-Compliant)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "env-cell", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"\n", + "CompilerOptimizationEnv β€” OpenEnv-compatible RL Environment\n", + "OpenEnv Hackathon 2026 | Theme #2: Long-Horizon Planning\n", + "\n", + "Key improvements over naive baseline:\n", + " βœ“ Inherits from MCPEnvironment (OpenEnv API compliant)\n", + " βœ“ Dynamic crash_penalty scaled to reward range (not hardcoded -1000)\n", + " βœ“ Soft termination: invalid actions give 3 chances before episode ends\n", + " βœ“ No-op detection: penalises passes that change nothing\n", + " βœ“ Terminal bonus: rewards cumulative improvement, not just greedy steps\n", + " βœ“ StepResult / EpisodeStats dataclasses for clean interfaces\n", + " βœ“ Curriculum level support\n", + "\"\"\"\n", + "\n", + "import copy\n", + "import math\n", + "from dataclasses import dataclass, field\n", + "from typing import Any, Dict, List, Optional, Tuple\n", + "\n", + "\n", + "# ── Stub base class ───────────────────────────────────────────────────────────\n", + "# In production replace with: from openenv import MCPEnvironment\n", + "class MCPEnvironment:\n", + " \"\"\"Minimal stub. In production: `from openenv import MCPEnvironment`\"\"\"\n", + " def reset(self, *args, **kwargs): raise NotImplementedError\n", + " def step(self, *args, **kwargs): raise NotImplementedError\n", + " def state(self): raise NotImplementedError\n", + "\n", + "\n", + "# ── Data classes ──────────────────────────────────────────────────────────────\n", + "@dataclass\n", + "class StepResult:\n", + " observation: str\n", + " reward: float\n", + " done: bool\n", + " info: Dict[str, Any] = field(default_factory=dict)\n", + "\n", + "\n", + "@dataclass\n", + "class EpisodeStats:\n", + " steps_taken: int = 0\n", + " total_reward: float = 0.0\n", + " passes_applied: List[str] = field(default_factory=list)\n", + " invalid_actions: int = 0\n", + " no_ops: int = 0\n", + " baseline_cycles: int = 0\n", + " final_cycles: int = 0\n", + "\n", + " @property\n", + " def total_improvement_pct(self) -> float:\n", + " if self.baseline_cycles == 0:\n", + " return 0.0\n", + " return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0\n", + "\n", + "\n", + "# ── Core Environment ──────────────────────────────────────────────────────────\n", + "class CompilerOptimizationEnv(MCPEnvironment):\n", + " \"\"\"\n", + " RL environment where an LLM agent sequentially applies compiler\n", + " optimization passes to minimise CPU cycle count while preserving\n", + " program semantics.\n", + "\n", + " State : Pseudo-assembly representation of the current IR (text)\n", + " Actions: Named compiler passes (strings) from the `passes` dictionary\n", + " Reward : % cycle improvement per step βˆ’ time_tax, with terminal bonus\n", + " Done : Semantic violation | max steps reached | too many invalid actions\n", + " \"\"\"\n", + "\n", + " # Hyperparameters\n", + " TIME_TAX: float = 1.0\n", + " NO_OP_PENALTY: float = -2.0\n", + " INVALID_ACTION_PENALTY: float = -5.0\n", + " MAX_INVALID_ACTIONS: int = 3\n", + " TERMINAL_BONUS_SCALE: float = 0.5\n", + "\n", + " def __init__(\n", + " self,\n", + " role1_engine,\n", + " role3_passes: Dict[str, Any],\n", + " max_steps: int = 10,\n", + " curriculum_level: int = 1,\n", + " ):\n", + " self.engine = role1_engine\n", + " self.passes = role3_passes\n", + " self.max_steps = max_steps\n", + " self.curriculum_level = curriculum_level\n", + " self._valid_actions = frozenset(self.passes.keys())\n", + "\n", + " # Episode state\n", + " self._stats: Optional[EpisodeStats] = None\n", + " self.original_program = None\n", + " self.current_program = None\n", + " self.previous_cycles = 0\n", + " self._consecutive_invalid = 0\n", + "\n", + " # ── OpenEnv API ───────────────────────────────────────────────────────────\n", + " def reset(self, new_program_json: List[Dict]) -> str:\n", + " self.original_program = copy.deepcopy(new_program_json)\n", + " self.current_program = copy.deepcopy(new_program_json)\n", + " self.previous_cycles = self._safe_count_cycles(self.current_program)\n", + " self._consecutive_invalid = 0\n", + " self._stats = EpisodeStats(\n", + " baseline_cycles=self.previous_cycles,\n", + " final_cycles=self.previous_cycles,\n", + " )\n", + " return self.state()\n", + "\n", + " def state(self) -> str:\n", + " return self._program_to_pseudoasm(self.current_program)\n", + "\n", + " def step(self, action_string: str) -> StepResult:\n", + " assert self._stats is not None, \"Call reset() before step().\"\n", + " self._stats.steps_taken += 1\n", + "\n", + " if action_string not in self._valid_actions:\n", + " return self._handle_invalid_action(action_string)\n", + "\n", + " candidate_program = self.passes[action_string](\n", + " copy.deepcopy(self.current_program)\n", + " )\n", + "\n", + " is_valid = self.engine.verify_equivalence(self.original_program, candidate_program)\n", + " if not is_valid:\n", + " return self._handle_semantic_violation()\n", + "\n", + " new_cycles = self._safe_count_cycles(candidate_program)\n", + " reward, info = self._compute_reward(action_string, new_cycles)\n", + "\n", + " self.current_program = candidate_program\n", + " self.previous_cycles = new_cycles\n", + " self._stats.final_cycles = new_cycles\n", + " self._stats.total_reward += reward\n", + " self._stats.passes_applied.append(action_string)\n", + " self._consecutive_invalid = 0\n", + "\n", + " done = self._stats.steps_taken >= self.max_steps\n", + " if done:\n", + " terminal_bonus = self._terminal_bonus()\n", + " reward += terminal_bonus\n", + " info[\"terminal_bonus\"] = terminal_bonus\n", + " info[\"reason\"] = \"max_steps_reached\"\n", + " info[\"episode_stats\"] = self._episode_summary()\n", + "\n", + " return StepResult(self.state(), reward, done, info)\n", + "\n", + " # ── Reward logic ──────────────────────────────────────────────────────────\n", + " def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, Dict]:\n", + " info: Dict[str, Any] = {\"action\": action}\n", + " if self.previous_cycles == 0:\n", + " return -self.TIME_TAX, {**info, \"note\": \"zero_baseline\"}\n", + "\n", + " delta_pct = ((self.previous_cycles - new_cycles) / self.previous_cycles) * 100.0\n", + "\n", + " if new_cycles == self.previous_cycles:\n", + " reward = self.NO_OP_PENALTY\n", + " self._stats.no_ops += 1\n", + " info[\"no_op\"] = True\n", + " else:\n", + " reward = delta_pct - self.TIME_TAX\n", + " info[\"delta_pct\"] = round(delta_pct, 3)\n", + "\n", + " info[\"prev_cycles\"] = self.previous_cycles\n", + " info[\"new_cycles\"] = new_cycles\n", + " return reward, info\n", + "\n", + " def _terminal_bonus(self) -> float:\n", + " return max(0.0, self._stats.total_improvement_pct * self.TERMINAL_BONUS_SCALE)\n", + "\n", + " def _compute_crash_penalty(self) -> float:\n", + " # 2Γ— best possible episode reward β€” always catastrophic, never overwhelming\n", + " return -2.0 * (100.0 * self.max_steps)\n", + "\n", + " # ── Error handlers ────────────────────────────────────────────────────────\n", + " def _handle_invalid_action(self, action: str) -> StepResult:\n", + " self._consecutive_invalid += 1\n", + " self._stats.invalid_actions += 1\n", + " done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS\n", + " info = {\n", + " \"error\": f\"Unknown action: '{action}'\",\n", + " \"valid_actions\": sorted(self._valid_actions),\n", + " \"consecutive_invalid\": self._consecutive_invalid,\n", + " }\n", + " if done:\n", + " info[\"reason\"] = \"too_many_invalid_actions\"\n", + " info[\"episode_stats\"] = self._episode_summary()\n", + " return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info)\n", + "\n", + " def _handle_semantic_violation(self) -> StepResult:\n", + " return StepResult(\n", + " self.state(),\n", + " self._compute_crash_penalty(),\n", + " True,\n", + " {\n", + " \"error\": \"Semantic equivalence check FAILED.\",\n", + " \"reason\": \"semantic_violation\",\n", + " \"episode_stats\": self._episode_summary(),\n", + " },\n", + " )\n", + "\n", + " # ── State representation ──────────────────────────────────────────────────\n", + " @staticmethod\n", + " def _program_to_pseudoasm(program: List[Dict]) -> str:\n", + " if not program:\n", + " return \"; (empty program)\"\n", + " lines = []\n", + " for i, instr in enumerate(program):\n", + " op = instr.get(\"op\", \"NOP\")\n", + " args = instr.get(\"args\", [])\n", + " dest = instr.get(\"dest\")\n", + " typ = instr.get(\"type\", \"\")\n", + " arg_str = \", \".join(str(a) for a in args)\n", + " type_hint = f\":{typ}\" if typ else \"\"\n", + " if dest:\n", + " lines.append(f\" {i:>3}: {dest}{type_hint} = {op} {arg_str}\")\n", + " else:\n", + " lines.append(f\" {i:>3}: {op} {arg_str}\")\n", + " return \"\\n\".join(lines)\n", + "\n", + " # ── Utilities ─────────────────────────────────────────────────────────────\n", + " def _safe_count_cycles(self, program: List[Dict]) -> int:\n", + " return max(0, int(self.engine.execute_and_count_cycles(program)))\n", + "\n", + " def _episode_summary(self) -> Dict:\n", + " s = self._stats\n", + " return {\n", + " \"steps\": s.steps_taken,\n", + " \"total_reward\": round(s.total_reward, 3),\n", + " \"passes_applied\": s.passes_applied,\n", + " \"invalid_actions\": s.invalid_actions,\n", + " \"no_ops\": s.no_ops,\n", + " \"baseline_cycles\": s.baseline_cycles,\n", + " \"final_cycles\": s.final_cycles,\n", + " \"total_improvement_pct\": round(s.total_improvement_pct, 3),\n", + " }\n", + "\n", + " def available_actions(self) -> List[str]:\n", + " return sorted(self._valid_actions)\n", + "\n", + "\n", + "print(\"βœ… CompilerOptimizationEnv defined\")" + ] + }, + { + "cell_type": "markdown", + "id": "smoke-header", + "metadata": {}, + "source": [ + "## πŸ§ͺ Cell 3 β€” Smoke Test (No GPU / Real Engine Needed)\n", + "Validates the entire reward pipeline with a mock engine. Run this before spending compute." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "smoke-test-cell", + "metadata": {}, + "outputs": [], + "source": [ + "class MockEngine:\n", + " \"\"\"Stub engine: cycles = instruction count, all programs semantically valid.\"\"\"\n", + " def execute_and_count_cycles(self, program):\n", + " return len(program)\n", + "\n", + " def verify_equivalence(self, original, candidate):\n", + " return True\n", + "\n", + "\n", + "MOCK_PASSES = {\n", + " \"constant_folding\": lambda p: p[:-1] if len(p) > 1 else p,\n", + " \"dead_code_elimination\": lambda p: p[:-1] if len(p) > 2 else p,\n", + " \"loop_unrolling\": lambda p: p, # intentional no-op for testing\n", + "}\n", + "\n", + "SAMPLE_PROGRAM = [\n", + " {\"op\": \"const\", \"dest\": \"x\", \"args\": [\"5\"], \"type\": \"int\"},\n", + " {\"op\": \"const\", \"dest\": \"y\", \"args\": [\"3\"], \"type\": \"int\"},\n", + " {\"op\": \"add\", \"dest\": \"z\", \"args\": [\"x\", \"y\"], \"type\": \"int\"},\n", + " {\"op\": \"mul\", \"dest\": \"w\", \"args\": [\"z\", \"x\"], \"type\": \"int\"},\n", + " {\"op\": \"ret\", \"args\": [\"w\"]},\n", + "]\n", + "\n", + "engine = MockEngine()\n", + "env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=6)\n", + "obs = env.reset(SAMPLE_PROGRAM)\n", + "\n", + "print(\"═\" * 55)\n", + "print(\"SMOKE TEST\")\n", + "print(\"═\" * 55)\n", + "print(f\"Initial state:\\n{obs}\")\n", + "print(f\"\\nBaseline cycles : {env.previous_cycles}\")\n", + "print(f\"Available actions: {env.available_actions()}\")\n", + "print()\n", + "\n", + "actions_to_try = [\n", + " \"constant_folding\",\n", + " \"dead_code_elimination\",\n", + " \"loop_unrolling\", # no-op\n", + " \"hallucinated_pass\", # invalid β€” but won't kill episode yet\n", + " \"constant_folding\",\n", + " \"dead_code_elimination\",\n", + "]\n", + "\n", + "for action in actions_to_try:\n", + " result = env.step(action)\n", + " tag = \"βœ—\" if result.reward < 0 else \"βœ“\"\n", + " print(f\"{tag} '{action}'\")\n", + " print(f\" reward={result.reward:+.2f} done={result.done}\")\n", + " relevant = {k: v for k, v in result.info.items()\n", + " if k in (\"delta_pct\", \"error\", \"no_op\", \"reason\",\n", + " \"terminal_bonus\", \"episode_stats\")}\n", + " if relevant:\n", + " print(f\" info: {relevant}\")\n", + " print()\n", + " if result.done:\n", + " break\n", + "\n", + "print(\"βœ… Smoke test passed\")" + ] + }, + { + "cell_type": "markdown", + "id": "reward-plot-header", + "metadata": {}, + "source": [ + "## πŸ“Š Cell 4 β€” Visualise Reward Across a Mock Episode" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "reward-plot-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import matplotlib.ticker as ticker\n", + "\n", + "# Run a full episode and collect data\n", + "env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10)\n", + "env.reset(SAMPLE_PROGRAM)\n", + "\n", + "sequence = [\n", + " \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n", + " \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n", + " \"dead_code_elimination\", \"loop_unrolling\", \"constant_folding\",\n", + " \"dead_code_elimination\",\n", + "]\n", + "\n", + "rewards, cycle_counts, actions_log = [], [], []\n", + "for act in sequence:\n", + " r = env.step(act)\n", + " rewards.append(r.reward)\n", + " cycle_counts.append(env.previous_cycles)\n", + " actions_log.append(act)\n", + " if r.done:\n", + " break\n", + "\n", + "steps = list(range(1, len(rewards) + 1))\n", + "\n", + "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n", + "fig.suptitle(\"Compiler Optimization Episode β€” Mock Engine\", fontsize=14, fontweight=\"bold\")\n", + "\n", + "# Reward per step\n", + "colors = [\"#2ecc71\" if r >= 0 else \"#e74c3c\" for r in rewards]\n", + "ax1.bar(steps, rewards, color=colors, edgecolor=\"white\", linewidth=0.5)\n", + "ax1.axhline(0, color=\"grey\", linewidth=0.8, linestyle=\"--\")\n", + "ax1.set_ylabel(\"Reward\")\n", + "ax1.set_title(\"Reward per Step (green = positive, red = negative)\")\n", + "ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter(\"%.1f\"))\n", + "\n", + "# Cycle count over time\n", + "ax2.plot(steps, cycle_counts, marker=\"o\", color=\"#3498db\", linewidth=2, markersize=6)\n", + "ax2.set_xlabel(\"Step\")\n", + "ax2.set_ylabel(\"CPU Cycles\")\n", + "ax2.set_title(\"CPU Cycle Count Over Episode (lower = better)\")\n", + "ax2.set_xticks(steps)\n", + "ax2.set_xticklabels(\n", + " [a.replace(\"_\", \"\\n\") for a in actions_log],\n", + " fontsize=7,\n", + ")\n", + "\n", + "plt.tight_layout()\n", + "plt.savefig(\"episode_reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(\"πŸ“ˆ Plot saved as episode_reward_curve.png\")" + ] + }, + { + "cell_type": "markdown", + "id": "model-header", + "metadata": {}, + "source": [ + "## πŸ€– Cell 5 β€” Load Model with Unsloth (QLoRA 4-bit)\n", + "> **Requires T4 GPU.** Skip to Cell 9 if you only want to test the environment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "model-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from unsloth import FastLanguageModel\n", + "\n", + "MODEL_NAME = \"unsloth/Qwen2.5-3B-Instruct\" # swap to 7B if VRAM allows\n", + "MAX_SEQ_LEN = 1024\n", + "LORA_RANK = 16\n", + "\n", + "model, tokenizer = FastLanguageModel.from_pretrained(\n", + " model_name = MODEL_NAME,\n", + " max_seq_length = MAX_SEQ_LEN,\n", + " dtype = None, # auto bf16/fp16\n", + " load_in_4bit = True,\n", + ")\n", + "\n", + "model = FastLanguageModel.get_peft_model(\n", + " model,\n", + " r = LORA_RANK,\n", + " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n", + " \"gate_proj\", \"up_proj\", \"down_proj\"],\n", + " lora_alpha = LORA_RANK * 2,\n", + " lora_dropout = 0.0,\n", + " bias = \"none\",\n", + " use_gradient_checkpointing = \"unsloth\",\n", + " random_state = 42,\n", + ")\n", + "\n", + "print(f\"βœ… Loaded {MODEL_NAME} with QLoRA rank={LORA_RANK}\")\n", + "print(f\" GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")" + ] + }, + { + "cell_type": "markdown", + "id": "prompt-header", + "metadata": {}, + "source": [ + "## πŸ’¬ Cell 6 β€” System Prompt & Dataset Builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "prompt-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from datasets import Dataset\n", + "\n", + "\n", + "def build_system_prompt(env: CompilerOptimizationEnv) -> str:\n", + " actions = \"\\n\".join(f\" - {a}\" for a in env.available_actions())\n", + " return (\n", + " \"You are a compiler optimization agent. Your goal is to reduce \"\n", + " \"CPU cycle count by applying optimization passes to the program below.\\n\\n\"\n", + " f\"Available actions (respond with EXACTLY one per turn):\\n{actions}\\n\"\n", + " \" - done (stop early if no further improvement is possible)\\n\\n\"\n", + " \"Rules:\\n\"\n", + " \" β€’ Output only the action name. No explanation, no markdown.\\n\"\n", + " \" β€’ Do not invent actions not listed above.\\n\"\n", + " \" β€’ Repeating a pass that does nothing wastes a step.\\n\"\n", + " )\n", + "\n", + "\n", + "def build_dataset(\n", + " programs: list,\n", + " engine,\n", + " passes: dict,\n", + ") -> Dataset:\n", + " \"\"\"\n", + " Each row = one episode's initial state.\n", + " GRPO samples K completions (actions) per row to estimate group-relative advantage.\n", + " \"\"\"\n", + " env = CompilerOptimizationEnv(engine, passes, max_steps=10)\n", + " system_prompt = build_system_prompt(env)\n", + "\n", + " rows = []\n", + " for prog in programs:\n", + " obs = env.reset(prog)\n", + " prompt = [\n", + " {\"role\": \"system\", \"content\": system_prompt},\n", + " {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n", + " ]\n", + " rows.append({\"prompt\": prompt, \"program_json\": json.dumps(prog)})\n", + "\n", + " return Dataset.from_list(rows)\n", + "\n", + "\n", + "# ── Demo: build dataset from mock programs ────────────────────────────────────\n", + "def make_mock_programs(n: int = 20) -> list:\n", + " \"\"\"Generate N random mock IR programs for demo purposes.\"\"\"\n", + " import random\n", + " ops = [\"add\", \"mul\", \"sub\", \"const\", \"load\"]\n", + " progs = []\n", + " for _ in range(n):\n", + " length = random.randint(4, 12)\n", + " prog = [\n", + " {\"op\": random.choice(ops),\n", + " \"dest\": f\"v{i}\",\n", + " \"args\": [f\"v{max(0,i-1)}\"],\n", + " \"type\": \"int\"}\n", + " for i in range(length)\n", + " ]\n", + " prog.append({\"op\": \"ret\", \"args\": [f\"v{length-1}\"]})\n", + " progs.append(prog)\n", + " return progs\n", + "\n", + "\n", + "mock_programs = make_mock_programs(n=30)\n", + "train_dataset = build_dataset(mock_programs, MockEngine(), MOCK_PASSES)\n", + "\n", + "print(f\"βœ… Dataset built: {len(train_dataset)} episodes\")\n", + "print(f\" Sample prompt keys: {list(train_dataset[0].keys())}\")" + ] + }, + { + "cell_type": "markdown", + "id": "reward-fn-header", + "metadata": {}, + "source": [ + "## 🎯 Cell 7 β€” Reward Function Factory (for GRPO)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "reward-fn-cell", + "metadata": {}, + "outputs": [], + "source": [ + "import re\n", + "from typing import Any\n", + "\n", + "\n", + "def make_reward_fn(engine, passes, max_steps=10):\n", + " \"\"\"\n", + " Returns a reward function compatible with TRL's GRPOTrainer.\n", + " Each GRPO rollout gets its own fresh env instance β€” no state leakage.\n", + "\n", + " Action parsing strips punctuation the model might add\n", + " (e.g. 'constant_folding.' β†’ 'constant_folding').\n", + " \"\"\"\n", + " def reward_fn(prompts, completions, **kwargs):\n", + " programs = kwargs.get(\"program_json\", [None] * len(completions))\n", + " rewards = []\n", + "\n", + " for completion, prog_json in zip(completions, programs):\n", + " # Parse action from model output\n", + " raw = completion if isinstance(completion, str) else completion[0][\"content\"]\n", + " action = raw.strip().lower().split()[0] if raw.strip() else \"__invalid__\"\n", + " action = re.sub(r\"[^a-z0-9_]\", \"\", action) # strip punctuation\n", + "\n", + " # Fresh env per rollout\n", + " env = CompilerOptimizationEnv(\n", + " role1_engine = engine,\n", + " role3_passes = passes,\n", + " max_steps = max_steps,\n", + " )\n", + " program = json.loads(prog_json) if prog_json else []\n", + " env.reset(program)\n", + "\n", + " if action == \"done\":\n", + " rewards.append(0.0) # neutral stop\n", + " else:\n", + " result = env.step(action)\n", + " rewards.append(result.reward)\n", + "\n", + " return rewards\n", + "\n", + " return reward_fn\n", + "\n", + "\n", + "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=10)\n", + "print(\"βœ… Reward function factory ready\")\n", + "\n", + "# Quick sanity check\n", + "test_completions = [\"constant_folding\", \"hallucinated_pass\", \"loop_unrolling\"]\n", + "test_programs = [json.dumps(SAMPLE_PROGRAM)] * 3\n", + "test_rewards = reward_fn(\n", + " prompts = [\"\"] * 3,\n", + " completions = test_completions,\n", + " program_json = test_programs,\n", + ")\n", + "print(\"\\nReward sanity check:\")\n", + "for act, rew in zip(test_completions, test_rewards):\n", + " print(f\" '{act}' β†’ {rew:+.2f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "trainer-header", + "metadata": {}, + "source": [ + "## πŸš€ Cell 8 β€” GRPO Trainer Config & Training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "trainer-cell", + "metadata": {}, + "outputs": [], + "source": [ + "from trl import GRPOConfig, GRPOTrainer\n", + "\n", + "# Optional W&B β€” comment out if not using\n", + "try:\n", + " import wandb\n", + " wandb.init(project=\"openenv-compiler-opt\", name=\"grpo-qwen2.5-3b\")\n", + " REPORT_TO = \"wandb\"\n", + "except Exception:\n", + " REPORT_TO = \"none\"\n", + "\n", + "\n", + "grpo_config = GRPOConfig(\n", + " # ── Generation ────────────────────────────────────────────────────────\n", + " num_generations = 4, # K rollouts per prompt for group-relative advantage\n", + " max_new_tokens = 16, # Actions are 1 word; don't waste context\n", + " temperature = 0.9,\n", + " top_p = 0.95,\n", + "\n", + " # ── Optimisation ──────────────────────────────────────────────────────\n", + " learning_rate = 5e-6,\n", + " per_device_train_batch_size = 2,\n", + " gradient_accumulation_steps = 4, # effective batch = 8\n", + " num_train_epochs = 3,\n", + " max_grad_norm = 0.5,\n", + "\n", + " # ── GRPO-specific ─────────────────────────────────────────────────────\n", + " beta = 0.04, # KL penalty; keeps policy near reference\n", + "\n", + " # ── Logging / checkpointing ───────────────────────────────────────────\n", + " output_dir = \"./grpo_compiler_checkpoints\",\n", + " logging_steps = 10,\n", + " save_steps = 100,\n", + " report_to = REPORT_TO,\n", + "\n", + " # ── Reproducibility ───────────────────────────────────────────────────\n", + " seed = 42,\n", + ")\n", + "\n", + "trainer = GRPOTrainer(\n", + " model = model,\n", + " tokenizer = tokenizer,\n", + " config = grpo_config,\n", + " train_dataset = train_dataset,\n", + " reward_funcs = reward_fn,\n", + ")\n", + "\n", + "print(\"βœ… Trainer configured\")\n", + "print(f\" num_generations (K) = {grpo_config.num_generations}\")\n", + "print(f\" effective batch size = \"\n", + " f\"{grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n", + "print(f\" KL beta = {grpo_config.beta}\")\n", + "print()\n", + "print(\"Starting training... (this will take a while on T4)\")\n", + "trainer.train()" + ] + }, + { + "cell_type": "markdown", + "id": "save-header", + "metadata": {}, + "source": [ + "## πŸ’Ύ Cell 9 β€” Save Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "save-cell", + "metadata": {}, + "outputs": [], + "source": [ + "SAVE_PATH = \"./grpo_compiler_final\"\n", + "\n", + "model.save_pretrained(SAVE_PATH)\n", + "tokenizer.save_pretrained(SAVE_PATH)\n", + "\n", + "print(f\"βœ… Model saved to {SAVE_PATH}\")\n", + "\n", + "# Optional: push to HuggingFace Hub\n", + "# model.push_to_hub(\"your-hf-username/compiler-opt-grpo\")\n", + "# tokenizer.push_to_hub(\"your-hf-username/compiler-opt-grpo\")" + ] + }, + { + "cell_type": "markdown", + "id": "curriculum-header", + "metadata": {}, + "source": [ + "## πŸ“ˆ Cell 10 β€” Curriculum Callback & Reward Tracking" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "curriculum-cell", + "metadata": {}, + "outputs": [], + "source": [ + "class CurriculumCallback:\n", + " \"\"\"\n", + " Tracks rolling mean reward and promotes curriculum level\n", + " when the agent has mastered the current difficulty.\n", + "\n", + " Usage: call .record(reward) after every episode.\n", + " Read .level to get current difficulty (1=easy, 2=medium, 3=hard).\n", + " \"\"\"\n", + " def __init__(self, reward_threshold: float = 5.0, window: int = 50):\n", + " self.threshold = reward_threshold\n", + " self.window = window\n", + " self._history = []\n", + " self.level = 1\n", + " self._promotions = []\n", + "\n", + " def record(self, reward: float, step: int = None):\n", + " self._history.append(reward)\n", + " if len(self._history) >= self.window:\n", + " mean = sum(self._history[-self.window:]) / self.window\n", + " if mean >= self.threshold and self.level < 3:\n", + " self.level += 1\n", + " self._promotions.append((step or len(self._history), self.level))\n", + " print(f\"[Curriculum] β–² Promoted to level {self.level} \"\n", + " f\"(rolling mean={mean:.2f})\")\n", + "\n", + " def plot(self):\n", + " import matplotlib.pyplot as plt\n", + " import numpy as np\n", + "\n", + " history = self._history\n", + " steps = list(range(len(history)))\n", + " window = self.window\n", + " rolling = [\n", + " sum(history[max(0,i-window):i+1]) / min(i+1, window)\n", + " for i in steps\n", + " ]\n", + "\n", + " fig, ax = plt.subplots(figsize=(10, 4))\n", + " ax.plot(steps, history, alpha=0.3, color=\"#3498db\", label=\"Episode reward\")\n", + " ax.plot(steps, rolling, color=\"#e74c3c\", linewidth=2,\n", + " label=f\"Rolling mean (w={window})\")\n", + " ax.axhline(self.threshold, linestyle=\"--\", color=\"grey\",\n", + " linewidth=1, label=f\"Promotion threshold ({self.threshold})\")\n", + " for step, level in self._promotions:\n", + " ax.axvline(step, color=\"green\", linewidth=1.5, linestyle=\":\")\n", + " ax.text(step, ax.get_ylim()[1]*0.9, f\" L{level}\",\n", + " color=\"green\", fontsize=9)\n", + " ax.set_xlabel(\"Episode\")\n", + " ax.set_ylabel(\"Reward\")\n", + " ax.set_title(\"Training Reward + Curriculum Progression\")\n", + " ax.legend()\n", + " plt.tight_layout()\n", + " plt.savefig(\"curriculum_reward_curve.png\", dpi=150)\n", + " plt.show()\n", + " print(\"πŸ“ˆ Saved curriculum_reward_curve.png\")\n", + "\n", + "\n", + "# ── Demo: simulate 200 episodes of improving reward ──────────────────────────\n", + "import random\n", + "cb = CurriculumCallback(reward_threshold=5.0, window=50)\n", + "for ep in range(200):\n", + " # Simulate reward slowly improving\n", + " synthetic_reward = -5 + ep * 0.08 + random.gauss(0, 2)\n", + " cb.record(synthetic_reward, step=ep)\n", + "\n", + "cb.plot()" + ] + }, + { + "cell_type": "markdown", + "id": "inference-header", + "metadata": {}, + "source": [ + "## πŸ” Cell 11 β€” Inference: Before vs After Training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "inference-cell", + "metadata": {}, + "outputs": [], + "source": [ + "def run_inference_episode(model, tokenizer, engine, passes, program, max_steps=10):\n", + " \"\"\"Run a greedy inference episode and return the action sequence + total improvement.\"\"\"\n", + " from transformers import TextStreamer\n", + " FastLanguageModel.for_inference(model)\n", + "\n", + " env = CompilerOptimizationEnv(engine, passes, max_steps=max_steps)\n", + " obs = env.reset(program)\n", + " system_prompt = build_system_prompt(env)\n", + "\n", + " actions_chosen, rewards_earned = [], []\n", + " done = False\n", + "\n", + " print(f\"\\nBaseline cycles: {env.previous_cycles}\")\n", + " print(f\"Initial state:\\n{obs}\\n\")\n", + "\n", + " while not done:\n", + " messages = [\n", + " {\"role\": \"system\", \"content\": system_prompt},\n", + " {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n", + " ]\n", + " inputs = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=True,\n", + " add_generation_prompt=True,\n", + " return_tensors=\"pt\",\n", + " ).to(model.device)\n", + "\n", + " with torch.no_grad():\n", + " outputs = model.generate(\n", + " input_ids = inputs,\n", + " max_new_tokens = 16,\n", + " temperature = 0.1, # greedy-ish for inference\n", + " do_sample = True,\n", + " )\n", + "\n", + " raw_action = tokenizer.decode(\n", + " outputs[0][inputs.shape[-1]:], skip_special_tokens=True\n", + " ).strip()\n", + " action = re.sub(r\"[^a-z0-9_]\", \"\", raw_action.lower().split()[0])\n", + "\n", + " result = env.step(action)\n", + " actions_chosen.append(action)\n", + " rewards_earned.append(result.reward)\n", + " obs = result.observation\n", + " done = result.done\n", + "\n", + " print(f\"Step {len(actions_chosen)}: '{action}' β†’ reward={result.reward:+.2f}\")\n", + "\n", + " summary = env._episode_summary()\n", + " print(f\"\\n{'─'*40}\")\n", + " print(f\"Total improvement: {summary['total_improvement_pct']:.1f}%\")\n", + " print(f\"Final cycles: {summary['final_cycles']} (was {summary['baseline_cycles']})\")\n", + " return summary\n", + "\n", + "\n", + "# Uncomment after training:\n", + "# summary = run_inference_episode(\n", + "# model, tokenizer, MockEngine(), MOCK_PASSES, SAMPLE_PROGRAM\n", + "# )\n", + "\n", + "print(\"βœ… Inference cell ready. Uncomment the last block after training to run.\")" + ] + }, + { + "cell_type": "markdown", + "id": "tips-header", + "metadata": {}, + "source": [ + "---\n", + "## πŸ“ Notes & Tips\n", + "\n", + "| What | Why it matters |\n", + "|------|----------------|\n", + "| `num_generations=4` | GRPO needs Kβ‰₯2 rollouts per prompt to compute group-relative advantage. K=4 balances diversity vs. compute. |\n", + "| `beta=0.04` | KL penalty keeping policy close to reference. Too low β†’ mode collapse. Too high β†’ no learning. |\n", + "| `max_new_tokens=16` | Actions are one word. This prevents wasted computation and keeps the model from adding explanations. |\n", + "| Dynamic crash penalty | Computed as `βˆ’2 Γ— (100 Γ— max_steps)` so it always dominates the best possible episode reward without washing out all other reward signal. |\n", + "| Terminal bonus | Rewards cumulative improvement, not just greedy single-step gains. Critical for long-horizon tasks. |\n", + "| Soft invalid-action termination | 3 consecutive invalid actions β†’ end. Single mistakes don't kill the episode; the agent can recover. |\n", + "| `deepcopy` on all pass inputs | Role 3's passes mutate dicts in-place. Without this, `original_program` gets corrupted and the verifier fails spuriously. |\n", + "\n", + "**Next steps:**\n", + "- Replace `MockEngine` with Role 1's real engine\n", + "- Replace `MOCK_PASSES` with Role 3's real passes\n", + "- Push environment to HuggingFace Spaces: `openenv init && openenv deploy`\n", + "- Add W&B sweep to tune `beta`, `learning_rate`, `num_generations`" + ] + } + ] +}" diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/role2_deliverable3_training_loop (2) (1) (1).ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/role2_deliverable3_training_loop (2) (1) (1).ipynb new file mode 100644 index 0000000000000000000000000000000000000000..a6c1973d9daa43844c88f2cca3e0cc9c186314d3 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/role2_deliverable3_training_loop (2) (1) (1).ipynb @@ -0,0 +1,4047 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "hT5ZwZBzexRs" + }, + "source": [ + "# Role 2 β€” Deliverable 3: Unsloth / TRL GRPO Training Loop\n", + "\n", + "**Compatible with:**\n", + "```\n", + "trl==0.23.1 | unsloth==2026.4.8 | transformers==4.57.1\n", + "vllm==0.11.0 | torch==2.8.0+cu128 | GPU: Tesla T4 (CUDA 7.5)\n", + "```\n", + "\n", + "**Integration notes for teammates:**\n", + "- **Role 3** β†’ pass `curriculum_items` list into the training cell\n", + "- Reward function in this notebook is intentionally **format-only** (GRPO wiring sanity mode)\n", + "- `RewardCurveCallback` auto-saves `reward_loss_curve.png`\n", + "\n", + "**Key T4 fixes vs original notebook:**\n", + "- `fast_inference=True` removed β†’ was triggering vLLM which crashes on CUDA 7.5\n", + "- `vllm_enforce_eager` removed β†’ not a valid TRL 0.23.1 kwarg\n", + "- `use_vllm=False` enforced β†’ HF native generation used instead\n", + "- `per_device_train_batch_size = num_generations` set explicitly to silence Unsloth warnings" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "i5Ex9FwjexRu" + }, + "source": [ + "## Cell 1 β€” Install Dependencies" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "evmukG8EexRv", + "outputId": "88a2c882-9bec-4755-8539-c6568b17c301" + }, + "source": [ + "!pip install bitsandbytes --quiet\n", + "!pip install \\\n", + " \"trl[vllm]==0.23.1\" \\\n", + " \"vllm==0.11.0\" \\\n", + " \"transformers==4.57.1\" \\\n", + " \"unsloth @ git+https://github.com/unslothai/unsloth.git\" \\\n", + " \"unsloth_zoo\" \\\n", + " \"datasets\" \"accelerate\" \"peft\" \\\n", + " --quiet\n", + "\n", + "!pip uninstall torchcodec -y --quiet\n", + "\n", + "import subprocess\n", + "subprocess.run([\"ldconfig\", \"/usr/lib64-nvidia\"], capture_output=True)\n", + "subprocess.run([\"ldconfig\", \"/usr/local/cuda/lib64\"], capture_output=True)\n", + "\n", + "import bitsandbytes as bnb\n", + "print(\"βœ… bitsandbytes:\", bnb.__version__)" + ], + "execution_count": 1, + "outputs": [ + { + "output_type": "stream", + "text": [ + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 MB\u001b[0m \u001b[31m18.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25h Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n", + " Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n", + " Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.0/44.0 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m564.6/564.6 kB\u001b[0m \u001b[31m41.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m438.2/438.2 MB\u001b[0m \u001b[31m3.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m12.0/12.0 MB\u001b[0m \u001b[31m124.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m180.0/180.0 kB\u001b[0m \u001b[31m20.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.5/45.5 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m111.0/111.0 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.4/45.4 kB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.9/3.9 MB\u001b[0m \u001b[31m107.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m110.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m887.9/887.9 MB\u001b[0m \u001b[31m1.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.0/4.0 MB\u001b[0m \u001b[31m111.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.6/8.6 MB\u001b[0m \u001b[31m111.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.2/117.2 MB\u001b[0m \u001b[31m8.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.7/8.7 MB\u001b[0m \u001b[31m112.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m322.4/322.4 MB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m155.6/155.6 MB\u001b[0m \u001b[31m6.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m421.9/421.9 kB\u001b[0m \u001b[31m42.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m506.8/506.8 kB\u001b[0m \u001b[31m49.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m114.2/114.2 kB\u001b[0m \u001b[31m13.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m566.4/566.4 kB\u001b[0m \u001b[31m49.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m15.0/15.0 MB\u001b[0m \u001b[31m107.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.5/6.5 MB\u001b[0m \u001b[31m135.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.0/3.0 MB\u001b[0m \u001b[31m114.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m48.9/48.9 MB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m73.8/73.8 MB\u001b[0m \u001b[31m8.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m23.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m33.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m388.0/388.0 kB\u001b[0m \u001b[31m13.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m288.2/288.2 kB\u001b[0m \u001b[31m12.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m51.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m225.0/225.0 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m180.7/180.7 kB\u001b[0m \u001b[31m8.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m71.6/71.6 kB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m185.2/185.2 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.4/42.4 MB\u001b[0m \u001b[31m12.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.5/79.5 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m331.1/331.1 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.0/8.0 MB\u001b[0m \u001b[31m32.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m819.4/819.4 kB\u001b[0m \u001b[31m23.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m959.8/959.8 kB\u001b[0m \u001b[31m23.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25h Building wheel for unsloth (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", + "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n", + "ipython 7.34.0 requires jedi>=0.16, which is not installed.\u001b[0m\u001b[31m\n", + "\u001b[0mβœ… bitsandbytes: 0.49.2\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8HfLIw8RexRv" + }, + "source": [ + "## Cell 2 β€” Version Check" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "5A3e24sGexRv", + "outputId": "9d72939b-4c7b-4ac2-97cf-ca30eecd01b0" + }, + "source": [ + "import unsloth # MUST be first import β€” patches trl/transformers\n", + "import vllm, trl, transformers\n", + "\n", + "print(\"vllm :\", vllm.__version__)\n", + "print(\"trl :\", trl.__version__)\n", + "print(\"transformers:\", transformers.__version__)\n", + "print(\"unsloth :\", unsloth.__version__)" + ], + "execution_count": 2, + "outputs": [ + { + "output_type": "stream", + "text": [ + "πŸ¦₯ Unsloth: Will patch your computer to enable 2x faster free finetuning.\n", + "INFO 04-25 13:31:38 [__init__.py:216] Automatically detected platform cuda.\n", + "ERROR 04-25 13:31:45 [fa_utils.py:57] Cannot use FA version 2 is not supported due to FA2 is only supported on devices with compute capability >= 8\n", + "πŸ¦₯ Unsloth Zoo will now patch everything to make training faster!\n", + "vllm : 0.11.0\n", + "trl : 0.23.1\n", + "transformers: 4.57.1\n", + "unsloth : 2026.4.8\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "R72vVp7JexRw" + }, + "source": [ + "## Cell 3 β€” Imports" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "HBZ9P_JiexRw", + "outputId": "bfb9ff12-f028-422f-8d4d-46664aec674d" + }, + "source": [ + "import os, json, textwrap, warnings\n", + "import torch\n", + "import matplotlib\n", + "matplotlib.use(\"Agg\") # headless β€” works in Colab & SSH\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from unsloth import FastLanguageModel\n", + "from datasets import Dataset\n", + "from transformers import TrainerCallback, TrainerControl, TrainerState\n", + "from trl import GRPOConfig, GRPOTrainer\n", + "\n", + "warnings.filterwarnings(\"ignore\")\n", + "print(\"βœ… Imports OK\")" + ], + "execution_count": 8, + "outputs": [ + { + "output_type": "stream", + "text": [ + "βœ… Imports OK\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "rRmA3CQPexRw" + }, + "source": [ + "## Cell 4 β€” Hyper-parameters (single source of truth)" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "spBue60aexRx", + "outputId": "d3aba521-933e-45ae-af56-7045e130ebd9" + }, + "source": [ + "CFG = dict(\n", + " model_name = \"unsloth/Qwen2.5-1.5B\", # swap to 3B if VRAM allows\n", + " max_seq_length = 512,\n", + " load_in_4bit = True,\n", + "\n", + " # LoRA\n", + " lora_r = 16,\n", + " lora_alpha = 16,\n", + " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n", + " \"gate_proj\", \"up_proj\", \"down_proj\"],\n", + "\n", + " # GRPO training\n", + " output_dir = \"grpo_outputs\",\n", + " learning_rate = 5e-6,\n", + " grad_accum = 1,\n", + " num_generations = 2, # completions sampled per prompt\n", + " max_new_tokens = 128, # keep short to stay within 512-token budget\n", + " temperature = 0.9,\n", + " max_steps = 300, # increase to 1000+ for real training\n", + " logging_steps = 5,\n", + " save_steps = 100,\n", + " seed = 3407,\n", + ")\n", + "\n", + "print(\"CFG loaded:\", CFG)" + ], + "execution_count": 9, + "outputs": [ + { + "output_type": "stream", + "text": [ + "CFG loaded: {'model_name': 'unsloth/Qwen2.5-1.5B', 'max_seq_length': 512, 'load_in_4bit': True, 'lora_r': 16, 'lora_alpha': 16, 'target_modules': ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'], 'output_dir': 'grpo_outputs', 'learning_rate': 5e-06, 'grad_accum': 1, 'num_generations': 2, 'max_new_tokens': 128, 'temperature': 0.9, 'max_steps': 300, 'logging_steps': 5, 'save_steps': 100, 'seed': 3407}\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "F60of_Ynjlo4", + "outputId": "19b4e13e-4b21-4b6e-88a2-3c896c8a7a2e" + }, + "source": [ + "!pip install -U bitsandbytes --quiet\n", + "!python -m bitsandbytes\n" + ], + "execution_count": 5, + "outputs": [ + { + "output_type": "stream", + "text": [ + "=================== bitsandbytes v0.49.2 ===================\n", + "Platform: Linux-6.6.113+-x86_64-with-glibc2.35\n", + " libc: glibc-2.35\n", + "Python: 3.12.13\n", + "PyTorch: 2.8.0+cu128\n", + " CUDA: 12.8\n", + " HIP: N/A\n", + " XPU: N/A\n", + "Related packages:\n", + " accelerate: 1.13.0\n", + " diffusers: 0.37.1\n", + " numpy: 2.0.2\n", + " pip: 24.1.2\n", + " peft: 0.18.1\n", + " safetensors: 0.7.0\n", + " transformers: 4.57.1\n", + " triton: 3.4.0\n", + " trl: 0.23.1\n", + "============================================================\n", + "PyTorch settings found: CUDA_VERSION=128, Highest Compute Capability: (7, 5).\n", + "Checking that the library is importable and CUDA is callable...\n", + "SUCCESS!\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "hL3xJ6QGjzyn", + "outputId": "30e131ae-873b-4947-cf4f-f7fd12a8c387" + }, + "source": [ + "import ctypes\n", + "ctypes.CDLL(\"/usr/lib64-nvidia/libcuda.so.1\")\n", + "\n", + "!ldconfig /usr/lib64-nvidia\n", + "!ldconfig /usr/local/cuda/lib64\n", + "\n", + "!pip install -U bitsandbytes --quiet\n", + "import importlib, bitsandbytes\n", + "importlib.reload(bitsandbytes)\n", + "print(\"βœ…\", bitsandbytes.__version__)" + ], + "execution_count": 6, + "outputs": [ + { + "output_type": "stream", + "text": [ + "/sbin/ldconfig.real: /usr/local/lib/libtcm.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc_proxy.so.2 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libhwloc.so.15 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc.so.2 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtcm_debug.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_opencl.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libumf.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_loader.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_5.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbb.so.12 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_0.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_level_zero_v2.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_level_zero.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtcm.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc_proxy.so.2 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libhwloc.so.15 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc.so.2 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtcm_debug.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_opencl.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libumf.so.1 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_loader.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_5.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbb.so.12 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_0.so.3 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_level_zero_v2.so.0 is not a symbolic link\n", + "\n", + "/sbin/ldconfig.real: /usr/local/lib/libur_adapter_level_zero.so.0 is not a symbolic link\n", + "\n", + "βœ… 0.49.2\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "QENQw4O6exRx" + }, + "source": [ + "## Cell 5 β€” Load Model + LoRA\n", + "\n", + "> **T4 fix:** `fast_inference=True` is intentionally omitted. That flag loads vLLM, which crashes on CUDA 7.5 with a tensor-size assertion (`expected size 2048==4096`). We use `use_vllm=False` in GRPOConfig instead." + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 464, + "referenced_widgets": [ + "a0e7598e41c14de1b47a056cad1d041e", + "0c21b9b4ee3a4d83b6032436ad61bfa8", + "fc89e3f5fe92417a8fa680d1c679aada", + "8393e3a15697431ca4df2c883a244c47", + "655d0a5f685c4ba389f7df0aeec1396f", + "370e8da3573d4bdaa3d23a042bd1190f", + "c3f10e9e96224ff7ad9b9bd689dcf93e", + "24d4b985e965409e9a01f7eb3ba3f23d", + "696f3bb68ab643a2b3afdc87a12debb8", + "1a63b047d97c4704837cd2d8b76692f9", + "f7aeef0949674a3fb0e7dcf70f07f8e4", + "fc0db98200be480cb63f9f4447433ee0", + "d6b1d2626c7c4c4097e363669c1e3bc7", + "df43d85bd2fd46dc82f746409a2a3d32", + "a0a15427a41c440fbcf2abcb6891aa23", + "8ac96d00937449da9846ea28c77addba", + "77a9a3a3564b4a42b516ef73f6b3eff2", + "49b8c93d9ff44b75b4dffc97bd99cfd7", + "a1533872921e4eeca8b76a81c1fa8dae", + "3d3575a4cd194727962f09f7d5b0c23d", + "15d47924927f4024848a6f6d8ee6acc4", + "3150ad407240488b8ea2db7778b7334c", + "de05b1ffd7fa42bda542ce3b5a662255", + "e578b95b72de4b62b5b0ac0e5a81c94a", + "a1ed2bfa9147481fa378df313f327426", + "6c0ba6aa76bd491896f0f7572bc6c5c8", + "e55e5e2f79d14fe48fef5e8f22eb89cf", + "893bbc7b1a6a489c857f658439055b02", + "a821eeb5758d472a9432b154dc4fdd0a", + "112fe371d54a443faa948d231ba3b778", + "c0ccdae3037b4aefbda9ae1c44f1798f", + "87745307bee0463ea7be429f51712b4f", + "b9c63458ed16499dafd94406d6cd9961", + "d5f3e631722641ed87ae7bba300acaaf", + "be94e8914efc4aa1887940fa075ff73a", + "9b4f23f724dd4d2fbe4e0bced4eed79d", + "f99bf961b8814ab78b01aec776bcd15e", + "93a7ad1ef9c44f26ba71bcce0a52fe2c", + "b96eb0324c064ddb99fd611fd16bd7c8", + "97b37500e2c549a0990db3ab687a203f", + "bd9a43c626dd4e7a9dff178911772961", + "5116b9c290ba4fac9e8dda430c6417d9", + "3a2c9d65c7f840aba8064f6358468edf", + "e331b7c239884d0d9fb476146a1ed616", + "e70ce7b46e22482aa3e793b3a5169aae", + "a7e1fcb6e01e453a8ae7f56f6faee951", + "74cba38a20344b80a6239fca5734a418", + "307d7f48695b417a814bf7ade5dccf80", + "e501375a0c4b4da09a9d686147aed523", + "57c8abc07c384b9da79a10fccd29e704", + "96c6e994c9764bbbb158dd945e0f6596", + "1b2373e813154598b83bdcb8e29e9d3e", + "d855da6d39b14ebe93134ef1f89c51d1", + "30ab4c100e3543ce8874ce82341ff399", + "caa9a03250c14102af84dcbed12b61f6", + "6156bd2c2d1144f08752fdf5c5085e14", + "b8260fc137e64f15a34a81f6dbacd2f3", + "f637e4e79fa2493c9489eecd965b8ce1", + "e94c42c578264e8894d99e1dd84249f3", + "ccbf6ed7c31b4a11bc31b4020ff7c0cf", + "4f6b9c2772654c29996cf6f98f9b3376", + "4963c7b82b11449e93df961873181e65", + "37c4302d7ba04d7ebf33222f30dc18b5", + "ba196ff67b0140918c62b20bc0a51cfc", + "a7706cbc52be4776bca2511b1f3ceed3", + "36ac26d37a904b44a01285f37e0fc685", + "b005cb13f6e14f98a6cdff7bc7406d6b", + "89a7db6c7b114a299a9a2ada9d9fc633", + "56a942afc0e84db0b9e65cb213a2270a", + "ea1ed850ad4d42949fdfed254e9e2ed1", + "50dd82810b07495bb85dc2bd82a2d9a9", + "44715b9483cb4d5c8626626d2c798a13", + "c37c1b43b04249bfb44ab5794f28c50f", + "c02a3f64ae464fddb84805823724447d", + "dce034aa439b4326a49bb81bd5f0f3f9", + "bf55337bded943a28115b236b5843a58", + "7a3d6b77271a4eebb277c3c25a39cb41", + "90e66331dac34ac2a6cc559f0ce7f270", + "8ef366851aab4de9a3fcf7f2f1e8300b", + "1068178decfb48c5a5a91f4501db96d0", + "bf7c1535d4e8478593b96e6ffa24f4db", + "b8c37b59d96542e8b2bd0d1818b67d7c", + "c7c8bc445e7d45f584754f80a5f7a91e", + "5d29795388a04241a4ef36e3200a39af", + "445ad499cd8d402e893715cfc3a22c1c", + "d4ed60dd9d0e4988986b49fc7c2369ac", + "b2071d06de0840758aca8184a50461cf", + "53601dd3a35c4ecea9e0d78881c1163b" + ] + }, + "id": "6ZQ4bBM6exRx", + "outputId": "a5a6a8d2-364b-452e-884a-0fb3a3582d62" + }, + "source": [ + "def load_model(cfg: dict):\n", + " print(\"β–Ά Loading base model …\")\n", + " model, tokenizer = FastLanguageModel.from_pretrained(\n", + " model_name = cfg[\"model_name\"],\n", + " max_seq_length = cfg[\"max_seq_length\"],\n", + " load_in_4bit = False, # bnb CUDA link broken on this runtime β€” 16bit fallback\n", + " dtype = torch.float16,\n", + " )\n", + "\n", + " print(\"β–Ά Attaching LoRA adapters …\")\n", + " model = FastLanguageModel.get_peft_model(\n", + " model,\n", + " r = cfg[\"lora_r\"],\n", + " target_modules = cfg[\"target_modules\"],\n", + " lora_alpha = cfg[\"lora_alpha\"],\n", + " use_gradient_checkpointing = \"unsloth\",\n", + " random_state = cfg[\"seed\"],\n", + " )\n", + " tokenizer.pad_token = tokenizer.eos_token\n", + " print(\"βœ… Model ready\")\n", + " return model, tokenizer\n", + "\n", + "\n", + "model, tokenizer = load_model(CFG)" + ], + "execution_count": 7, + "outputs": [ + { + "output_type": "stream", + "text": [ + "β–Ά Loading base model …\n", + "==((====))== Unsloth 2026.4.8: Fast Qwen2 patching. Transformers: 4.57.1. vLLM: 0.11.0.\n", + " \\\\ /| Tesla T4. Num GPUs = 1. Max memory: 14.563 GB. Platform: Linux.\n", + "O^O/ \\_/ \\ Torch: 2.8.0+cu128. CUDA: 7.5. CUDA Toolkit: 12.8. Triton: 3.4.0\n", + "\\ / Bfloat16 = FALSE. FA [Xformers = 0.0.32.post1. FA2 = False]\n", + " \"-____-\" Free license: http://github.com/unslothai/unsloth\n", + "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n" + ] + }, + { + "output_type": "display_data", + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a0e7598e41c14de1b47a056cad1d041e", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "model.safetensors: 0%| | 0.00/3.09G [00:00.\n", + "β–Ά Attaching LoRA adapters …\n" + ] + }, + { + "output_type": "stream", + "text": [ + "Unsloth 2026.4.8 patched 28 layers with 28 QKV layers, 28 O layers and 28 MLP layers.\n" + ] + }, + { + "output_type": "stream", + "text": [ + "βœ… Model ready\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "M_MNSnFXexRy" + }, + "source": [ + "## Cell 6 β€” Prompt Builder & Dataset\n", + "\n", + "Role 3 hands you a list of `{\"observation\": str, \"tac\": any, \"level\": int}` dicts. \n", + "Pass them to `build_dataset_from_curriculum()`. If Role 3 isn't ready yet, `_make_dummy_dataset()` is the fallback." + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "4TxF3BF7exRy", + "outputId": "27a8c4aa-9696-4c62-ddb8-2e4dc6117001" + }, + "source": [ + "import re\n", + "\n", + "SYSTEM_PROMPT = textwrap.dedent(\"\"\"\\\n", + " You are a compiler optimisation agent. You will receive a Toy-IR program\n", + " as pseudo-assembly (translated from Three-Address Code). Your job is to output\n", + " a JSON array of optimisation pass names to apply IN ORDER.\n", + "\n", + " Available passes: \\\"constant_folding\\\", \\\"dead_code_elimination\\\",\n", + " \\\"peephole_optimization\\\"\n", + "\n", + " Rules:\n", + " β€’ Output ONLY a valid JSON array, nothing else.\n", + " β€’ You may repeat passes.\n", + " β€’ Maximum 8 passes per sequence.\n", + "\n", + " Example output: [\\\"constant_folding\\\", \\\"dead_code_elimination\\\", \\\"peephole_optimization\\\"]\n", + "\"\"\")\n", + "\n", + "\n", + "class Deliverable2_Formatter:\n", + " @staticmethod\n", + " def translate_state(raw_json: list) -> str:\n", + " \"\"\"Translate raw JSON IR into compact pseudo-assembly.\"\"\"\n", + " if not isinstance(raw_json, list) or not raw_json:\n", + " return \"; (empty program β€” 0 instructions)\"\n", + "\n", + " pseudo_assembly = []\n", + " for i, instruction in enumerate(raw_json):\n", + " if not isinstance(instruction, dict):\n", + " pseudo_assembly.append(f\"{i}. NOP\")\n", + " continue\n", + "\n", + " op = str(instruction.get(\"op\", \"UNKNOWN\")).upper()\n", + " args = \", \".join(str(arg) for arg in instruction.get(\"args\", []))\n", + " dest = instruction.get(\"dest\", \"\")\n", + " if dest:\n", + " line = f\"{i}. {dest} = {op} {args}\".rstrip()\n", + " else:\n", + " line = f\"{i}. {op} {args}\".rstrip()\n", + " pseudo_assembly.append(line)\n", + "\n", + " return \"\\n\".join(pseudo_assembly)\n", + "\n", + " @staticmethod\n", + " def extract_action_array(llm_output: str) -> list:\n", + " \"\"\"Best-effort extraction of JSON pass arrays from noisy LLM output.\"\"\"\n", + " text = (llm_output or \"\").strip()\n", + " if not text:\n", + " raise ValueError(\"Invalid JSON format\")\n", + "\n", + " # Strategy 1: direct JSON\n", + " try:\n", + " parsed = json.loads(text)\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " pass\n", + "\n", + " # Strategy 2: markdown fence cleanup\n", + " cleaned = re.sub(r\"```(?:json)?\", \"\", text, flags=re.IGNORECASE).replace(\"```\", \"\").strip()\n", + " if cleaned != text:\n", + " try:\n", + " parsed = json.loads(cleaned)\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " pass\n", + "\n", + " # Strategy 3: bracket extraction\n", + " match = re.search(r\"\\[.*?\\]\", text, re.DOTALL)\n", + " if match:\n", + " candidate = match.group(0)\n", + " try:\n", + " parsed = json.loads(candidate)\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " try:\n", + " parsed = json.loads(candidate.replace(\"'\", '\"'))\n", + " if isinstance(parsed, list):\n", + " return parsed\n", + " except json.JSONDecodeError:\n", + " pass\n", + "\n", + " raise ValueError(\"Invalid JSON format\")\n", + "\n", + "\n", + "def build_prompt(tac_or_ir) -> str:\n", + " \"\"\"Convert either TAC string or raw IR list into a chat prompt.\"\"\"\n", + " tac_string = (\n", + " Deliverable2_Formatter.translate_state(tac_or_ir)\n", + " if isinstance(tac_or_ir, list)\n", + " else str(tac_or_ir)\n", + " )\n", + " return (\n", + " f\"<|im_start|>system\\n{SYSTEM_PROMPT}<|im_end|>\\n\"\n", + " f\"<|im_start|>user\\n{tac_string}<|im_end|>\\n\"\n", + " f\"<|im_start|>assistant\\n\"\n", + " )\n", + "\n", + "\n", + "def build_dataset_from_curriculum(curriculum_items: list) -> Dataset:\n", + " \"\"\"\n", + " Convert Role 3 curriculum output into a HF Dataset.\n", + " Supports either observation strings or raw TAC/IR lists.\n", + " \"\"\"\n", + " prompts = []\n", + " tac_payloads = []\n", + "\n", + " for item in curriculum_items:\n", + " tac_obj = item.get(\"tac\", item.get(\"observation\", \"\"))\n", + " if isinstance(tac_obj, list):\n", + " prompts.append(build_prompt(tac_obj))\n", + " else:\n", + " obs = item.get(\"observation\", tac_obj)\n", + " prompts.append(build_prompt(obs))\n", + " tac_payloads.append(json.dumps(tac_obj))\n", + "\n", + " return Dataset.from_dict({\"prompt\": prompts, \"tac_raw\": tac_payloads})\n", + "\n", + "\n", + "def _make_dummy_dataset() -> Dataset:\n", + " \"\"\"Fallback dataset β€” self-contained, no Role 3 dependency.\"\"\"\n", + " dummy_tacs = [\n", + " [\n", + " {\"op\": \"const\", \"dest\": \"t0\", \"args\": [\"3\"]},\n", + " {\"op\": \"const\", \"dest\": \"t1\", \"args\": [\"4\"]},\n", + " {\"op\": \"add\", \"dest\": \"t2\", \"args\": [\"t0\", \"t1\"]},\n", + " {\"op\": \"ret\", \"args\": [\"t2\"]},\n", + " ],\n", + " [\n", + " {\"op\": \"const\", \"dest\": \"a\", \"args\": [\"2\"]},\n", + " {\"op\": \"const\", \"dest\": \"b\", \"args\": [\"2\"]},\n", + " {\"op\": \"mul\", \"dest\": \"c\", \"args\": [\"a\", \"b\"]},\n", + " {\"op\": \"ret\", \"args\": [\"c\"]},\n", + " ],\n", + " ]\n", + " items = [{\"observation\": Deliverable2_Formatter.translate_state(t), \"tac\": t} for t in dummy_tacs]\n", + " return build_dataset_from_curriculum(items)\n", + "\n", + "\n", + "print(\"βœ… Prompt builder and dataset functions defined\")" + ], + "execution_count": 10, + "outputs": [ + { + "output_type": "stream", + "text": [ + "βœ… Prompt builder and dataset functions defined\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1458Ig0-exRy" + }, + "source": [ + "## Cell 7 β€” Reward Function\n", + "\n", + "This notebook is GRPO wiring only.\n", + "\n", + "The reward below scores output format quality (valid JSON pass list, bounded length, and diversity). It is not a semantic optimizer/verifier reward." + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "hHN3TiwaexRy", + "outputId": "11b0fee5-d901-461e-af2c-e52f68cc41a8" + }, + "source": [ + "VALID_PASSES = {\"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\"}\n", + "MAX_PASSES = 8\n", + "\n", + "\n", + "def _parse_pass_sequence(text: str):\n", + " \"\"\"Extract and validate pass list using Deliverable2_Formatter parser.\"\"\"\n", + " try:\n", + " passes = Deliverable2_Formatter.extract_action_array(text)\n", + " except ValueError:\n", + " return None\n", + "\n", + " if (\n", + " isinstance(passes, list)\n", + " and 0 < len(passes) <= MAX_PASSES\n", + " and all(isinstance(p, str) and p in VALID_PASSES for p in passes)\n", + " ):\n", + " return passes\n", + " return None\n", + "\n", + "\n", + "def make_reward_function():\n", + " \"\"\"Factory returning a TRL-compatible format-quality reward function.\"\"\"\n", + "\n", + " def reward_func(prompts: list, completions: list, **kwargs) -> list:\n", + " rewards = []\n", + " for _, completion in zip(prompts, completions):\n", + " passes = _parse_pass_sequence(completion)\n", + "\n", + " # Strict gate: invalid JSON array or unknown pass names.\n", + " if passes is None:\n", + " rewards.append(-1.0)\n", + " continue\n", + "\n", + " n = len(passes)\n", + " unique_ratio = len(set(passes)) / max(n, 1)\n", + " diversity = 0.25 * unique_ratio\n", + " length_calibration = 0.20 * min(n, 5) / 5.0\n", + " clean_format = completion.strip().startswith(\"[\") and completion.strip().endswith(\"]\")\n", + " format_bonus = 0.15 if clean_format else 0.0\n", + " reward = round(0.40 + diversity + length_calibration + format_bonus, 4)\n", + " rewards.append(reward)\n", + "\n", + " return rewards\n", + "\n", + " return reward_func\n", + "\n", + "\n", + "print(\"βœ… Reward function defined\")" + ], + "execution_count": 11, + "outputs": [ + { + "output_type": "stream", + "text": [ + "βœ… Reward function defined\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W0-iUouUexRy" + }, + "source": [ + "## Cell 8 β€” RewardCurveCallback\n", + "\n", + "Logs reward + loss at every `logging_steps`. Saves `reward_loss_curve.png` on `on_train_end`." + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "4HQW8VAXexRy", + "outputId": "c9e81817-a089-4284-e0f3-6c5767c209c5" + }, + "source": [ + "class RewardCurveCallback(TrainerCallback):\n", + " \"\"\"Logs mean reward and loss; saves dual-axis PNG for the storytelling deck.\"\"\"\n", + "\n", + " def __init__(self, output_dir: str):\n", + " self.output_dir = output_dir\n", + " self.steps: list = []\n", + " self.rewards: list = []\n", + " self.losses: list = []\n", + "\n", + " def on_log(self, args, state: TrainerState,\n", + " control: TrainerControl, logs: dict = None, **kwargs):\n", + " if logs is None:\n", + " return\n", + " step = state.global_step\n", + " reward = logs.get(\"reward\", logs.get(\"train/reward\", None))\n", + " loss = logs.get(\"loss\", logs.get(\"train/loss\", None))\n", + " if reward is not None:\n", + " self.steps.append(step)\n", + " self.rewards.append(float(reward))\n", + " self.losses.append(float(loss) if loss is not None else 0.0)\n", + " print(f\" step={step:4d} reward={float(reward):+.4f}\"\n", + " + (f\" loss={float(loss):.4f}\" if loss is not None else \"\"))\n", + "\n", + " def on_train_end(self, args, state, control, **kwargs):\n", + " if not self.steps:\n", + " return\n", + " fig, ax1 = plt.subplots(figsize=(10, 4))\n", + " ax1.plot(self.steps, self.rewards, \"b-o\", markersize=3, label=\"Mean Reward\")\n", + " ax1.set_xlabel(\"Step\")\n", + " ax1.set_ylabel(\"Reward\", color=\"b\")\n", + " ax1.tick_params(axis=\"y\", labelcolor=\"b\")\n", + " ax1.axhline(0, color=\"b\", linewidth=0.5, linestyle=\"--\")\n", + "\n", + " ax2 = ax1.twinx()\n", + " ax2.plot(self.steps, self.losses, \"r-s\", markersize=3, label=\"Loss\")\n", + " ax2.set_ylabel(\"Loss\", color=\"r\")\n", + " ax2.tick_params(axis=\"y\", labelcolor=\"r\")\n", + "\n", + " lines1, labels1 = ax1.get_legend_handles_labels()\n", + " lines2, labels2 = ax2.get_legend_handles_labels()\n", + " ax1.legend(lines1 + lines2, labels1 + labels2, loc=\"upper left\")\n", + " plt.title(\"GRPO Training β€” Reward & Loss Curves\")\n", + " plt.tight_layout()\n", + "\n", + " os.makedirs(self.output_dir, exist_ok=True)\n", + " path = os.path.join(self.output_dir, \"reward_loss_curve.png\")\n", + " plt.savefig(path, dpi=150)\n", + " plt.close()\n", + " print(f\"\\nβœ… Reward/loss curve saved β†’ {path}\")\n", + "\n", + "\n", + "print(\"βœ… RewardCurveCallback defined\")" + ], + "execution_count": 13, + "outputs": [ + { + "output_type": "stream", + "text": [ + "βœ… RewardCurveCallback defined\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "QAwDYqaVexRz" + }, + "source": [ + "## Cell 9 β€” GRPOConfig\n", + "\n", + "> **T4 fixes:** `use_vllm=False` (HF native generation). `vllm_enforce_eager` removed β€” not a valid TRL 0.23.1 kwarg. `per_device_train_batch_size` set equal to `num_generations` explicitly." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "HOZW2AZ4exRz" + }, + "source": [ + "from trl import GRPOConfig, GRPOTrainer\n", + "def build_grpo_config(cfg: dict) -> GRPOConfig:\n", + " return GRPOConfig(\n", + " output_dir = cfg[\"output_dir\"],\n", + " learning_rate = cfg[\"learning_rate\"],\n", + " per_device_train_batch_size = cfg[\"num_generations\"],\n", + " gradient_accumulation_steps = cfg[\"grad_accum\"],\n", + " num_generations = cfg[\"num_generations\"],\n", + " max_completion_length = cfg[\"max_new_tokens\"], # ← renamed in TRL 0.23.1\n", + " temperature = cfg[\"temperature\"],\n", + " max_steps = cfg[\"max_steps\"],\n", + " logging_steps = cfg[\"logging_steps\"],\n", + " save_steps = cfg[\"save_steps\"],\n", + " seed = cfg[\"seed\"],\n", + " report_to = \"none\",\n", + " use_vllm = False,\n", + " )" + ], + "execution_count": 14, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JQe_f_ukexRz" + }, + "source": [ + "## Cell 10 β€” Run Training\n", + "\n", + "**Sanity-check mode** (no teammates needed): leave `curriculum_items=None`. \n", + "**Full integration**: pass `curriculum_items` from Role 3's generator." + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 449 + }, + "id": "Pk5h045bexRz", + "outputId": "f6574776-e41f-4ab7-90d1-45d82e9b43e7" + }, + "source": [ + "# ── INTEGRATION: uncomment when Role 3 is ready ───────────────────────────────\n", + "# from role3_curriculum import generate_curriculum\n", + "# curriculum_items = generate_curriculum(n=200, max_level=5)\n", + "\n", + "# ── STANDALONE: dummy dataset fallback ────────────────────────────────────────\n", + "curriculum_items = None # ← set to Role 3 output when available\n", + "\n", + "# ── Build dataset ─────────────────────────────────────────────────────────────\n", + "if curriculum_items is not None:\n", + " print(f\"β–Ά Using Role 3 curriculum β€” {len(curriculum_items)} programs\")\n", + " dataset = build_dataset_from_curriculum(curriculum_items)\n", + "else:\n", + " print(\"β–Ά No curriculum provided β€” using built-in dummy dataset\")\n", + " dataset = _make_dummy_dataset()\n", + "\n", + "print(f\" Dataset size: {len(dataset)} examples\")\n", + "\n", + "# ── Build reward, config, trainer ─────────────────────────────────────────────\n", + "reward_fn = make_reward_function()\n", + "training_args = build_grpo_config(CFG)\n", + "reward_callback = RewardCurveCallback(CFG[\"output_dir\"])\n", + "\n", + "trainer = GRPOTrainer(\n", + " model = model,\n", + " processing_class = tokenizer,\n", + " reward_funcs = [reward_fn],\n", + " args = training_args,\n", + " train_dataset = dataset,\n", + " callbacks = [reward_callback],\n", + ")\n", + "\n", + "print(\"\\nπŸš€ Starting GRPO training …\\n\" + \"─\" * 50)\n", + "trainer.train()" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "text": [ + "β–Ά No curriculum provided β€” using built-in dummy dataset\n", + " Dataset size: 4 examples\n", + "\n", + "πŸš€ Starting GRPO training …\n", + "──────────────────────────────────────────────────\n" + ] + }, + { + "output_type": "stream", + "text": [ + "==((====))== Unsloth - 2x faster free finetuning | Num GPUs used = 1\n", + " \\\\ /| Num examples = 4 | Num Epochs = 75 | Total steps = 300\n", + "O^O/ \\_/ \\ Batch size per device = 2 | Gradient accumulation steps = 1\n", + "\\ / Data Parallel GPUs = 1 | Total batch size (2 x 1 x 1) = 2\n", + " \"-____-\" Trainable parameters = 18,464,768 of 1,562,179,072 (1.18% trained)\n" + ] + }, + { + "output_type": "display_data", + "data": { + "text/html": [ + "\n", + "
\n", + " \n", + " \n", + " [ 18/300 02:24 < 42:27, 0.11 it/s, Epoch 4.25/75]\n", + "
\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
StepTraining Lossrewardreward_stdcompletions / mean_lengthcompletions / min_lengthcompletions / max_lengthcompletions / clipped_ratiocompletions / mean_terminated_lengthcompletions / min_terminated_lengthcompletions / max_terminated_lengthklrewards / reward_func / meanrewards / reward_func / std
50.000000-1.0000000.000000106.400000104.400000108.4000000.8000004.0000002.0000006.0000000.000000-1.0000000.000000
100.000000-1.0000000.000000122.200000116.400000128.0000000.90000014.00000014.00000014.0000000.000000-1.0000000.000000
150.000000-0.8200000.254558117.800000107.600000128.0000000.80000030.80000030.80000030.8000000.000007-0.8200000.254558

" + ], + "text/plain": [ + "" + ] + } + }, + { + "output_type": "stream", + "text": [ + " step= 5 reward=-1.0000 loss=0.0000\n", + " step= 10 reward=-1.0000 loss=0.0000\n", + " step= 15 reward=-0.8200 loss=0.0000\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "c9yYMwOGexRz" + }, + "source": [ + "## Cell 11 β€” Save LoRA Adapter + Print Summary" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "D1iR6HmpexRz" + }, + "source": [ + "adapter_path = os.path.join(CFG[\"output_dir\"], \"lora_adapter\")\n", + "model.save_pretrained(adapter_path)\n", + "tokenizer.save_pretrained(adapter_path)\n", + "print(f\"βœ… LoRA adapter saved β†’ {adapter_path}\")\n", + "\n", + "if reward_callback.steps:\n", + " best = max(reward_callback.rewards)\n", + " final = reward_callback.rewards[-1]\n", + " print(f\"\\n{'='*50}\")\n", + " print(f\" PIPELINE STABLE!\")\n", + " print(f\" Best reward : {best:+.4f}\")\n", + " print(f\" Final reward: {final:+.4f}\")\n", + " print(f\" Total steps : {reward_callback.steps[-1]}\")\n", + " print(f\"{'='*50}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oczg4iVlexR0" + }, + "source": [ + "## Cell 12 β€” Display Reward/Loss Curve\n", + "\n", + "The PNG was auto-saved by `RewardCurveCallback`. This cell renders it inline." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "_2SdVxZ_exR0" + }, + "source": [ + "from IPython.display import Image, display\n", + "import os\n", + "\n", + "curve_path = os.path.join(CFG[\"output_dir\"], \"reward_loss_curve.png\")\n", + "if os.path.exists(curve_path):\n", + " display(Image(curve_path))\n", + "else:\n", + " print(\"Curve not saved yet β€” did training complete?\")" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "0c21b9b4ee3a4d83b6032436ad61bfa8": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_370e8da3573d4bdaa3d23a042bd1190f", + "placeholder": "​", + "style": "IPY_MODEL_c3f10e9e96224ff7ad9b9bd689dcf93e", + "value": "model.safetensors: 100%" + } + }, + "1068178decfb48c5a5a91f4501db96d0": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_445ad499cd8d402e893715cfc3a22c1c", + "max": 11421896, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_d4ed60dd9d0e4988986b49fc7c2369ac", + "value": 11421896 + } + }, + "112fe371d54a443faa948d231ba3b778": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": "20px" + } + }, + "15d47924927f4024848a6f6d8ee6acc4": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "1a63b047d97c4704837cd2d8b76692f9": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "1b2373e813154598b83bdcb8e29e9d3e": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": "20px" + } + }, + "24d4b985e965409e9a01f7eb3ba3f23d": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "307d7f48695b417a814bf7ade5dccf80": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_30ab4c100e3543ce8874ce82341ff399", + "placeholder": "​", + "style": "IPY_MODEL_caa9a03250c14102af84dcbed12b61f6", + "value": " 1.67M/? [00:00<00:00, 47.0MB/s]" + } + }, + "30ab4c100e3543ce8874ce82341ff399": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "3150ad407240488b8ea2db7778b7334c": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "36ac26d37a904b44a01285f37e0fc685": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "370e8da3573d4bdaa3d23a042bd1190f": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "37c4302d7ba04d7ebf33222f30dc18b5": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "3a2c9d65c7f840aba8064f6358468edf": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "3d3575a4cd194727962f09f7d5b0c23d": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "445ad499cd8d402e893715cfc3a22c1c": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "44715b9483cb4d5c8626626d2c798a13": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "4963c7b82b11449e93df961873181e65": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "49b8c93d9ff44b75b4dffc97bd99cfd7": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "4f6b9c2772654c29996cf6f98f9b3376": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "50dd82810b07495bb85dc2bd82a2d9a9": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "5116b9c290ba4fac9e8dda430c6417d9": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "53601dd3a35c4ecea9e0d78881c1163b": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "56a942afc0e84db0b9e65cb213a2270a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_c02a3f64ae464fddb84805823724447d", + "max": 617, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_dce034aa439b4326a49bb81bd5f0f3f9", + "value": 617 + } + }, + "57c8abc07c384b9da79a10fccd29e704": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "5d29795388a04241a4ef36e3200a39af": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "6156bd2c2d1144f08752fdf5c5085e14": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_b8260fc137e64f15a34a81f6dbacd2f3", + "IPY_MODEL_f637e4e79fa2493c9489eecd965b8ce1", + "IPY_MODEL_e94c42c578264e8894d99e1dd84249f3" + ], + "layout": "IPY_MODEL_ccbf6ed7c31b4a11bc31b4020ff7c0cf" + } + }, + "655d0a5f685c4ba389f7df0aeec1396f": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "696f3bb68ab643a2b3afdc87a12debb8": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "6c0ba6aa76bd491896f0f7572bc6c5c8": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_87745307bee0463ea7be429f51712b4f", + "placeholder": "​", + "style": "IPY_MODEL_b9c63458ed16499dafd94406d6cd9961", + "value": " 4.71k/? [00:00<00:00, 298kB/s]" + } + }, + "74cba38a20344b80a6239fca5734a418": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_1b2373e813154598b83bdcb8e29e9d3e", + "max": 1, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_d855da6d39b14ebe93134ef1f89c51d1", + "value": 1 + } + }, + "77a9a3a3564b4a42b516ef73f6b3eff2": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "7a3d6b77271a4eebb277c3c25a39cb41": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "8393e3a15697431ca4df2c883a244c47": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_1a63b047d97c4704837cd2d8b76692f9", + "placeholder": "​", + "style": "IPY_MODEL_f7aeef0949674a3fb0e7dcf70f07f8e4", + "value": " 3.09G/3.09G [00:27<00:00, 193MB/s]" + } + }, + "87745307bee0463ea7be429f51712b4f": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "893bbc7b1a6a489c857f658439055b02": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "89a7db6c7b114a299a9a2ada9d9fc633": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_44715b9483cb4d5c8626626d2c798a13", + "placeholder": "​", + "style": "IPY_MODEL_c37c1b43b04249bfb44ab5794f28c50f", + "value": "special_tokens_map.json: 100%" + } + }, + "8ac96d00937449da9846ea28c77addba": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "8ef366851aab4de9a3fcf7f2f1e8300b": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_c7c8bc445e7d45f584754f80a5f7a91e", + "placeholder": "​", + "style": "IPY_MODEL_5d29795388a04241a4ef36e3200a39af", + "value": "tokenizer.json: 100%" + } + }, + "90e66331dac34ac2a6cc559f0ce7f270": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_8ef366851aab4de9a3fcf7f2f1e8300b", + "IPY_MODEL_1068178decfb48c5a5a91f4501db96d0", + "IPY_MODEL_bf7c1535d4e8478593b96e6ffa24f4db" + ], + "layout": "IPY_MODEL_b8c37b59d96542e8b2bd0d1818b67d7c" + } + }, + "93a7ad1ef9c44f26ba71bcce0a52fe2c": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "96c6e994c9764bbbb158dd945e0f6596": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "97b37500e2c549a0990db3ab687a203f": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "9b4f23f724dd4d2fbe4e0bced4eed79d": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_bd9a43c626dd4e7a9dff178911772961", + "max": 1, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_5116b9c290ba4fac9e8dda430c6417d9", + "value": 1 + } + }, + "a0a15427a41c440fbcf2abcb6891aa23": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_15d47924927f4024848a6f6d8ee6acc4", + "placeholder": "​", + "style": "IPY_MODEL_3150ad407240488b8ea2db7778b7334c", + "value": " 171/171 [00:00<00:00, 12.7kB/s]" + } + }, + "a0e7598e41c14de1b47a056cad1d041e": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_0c21b9b4ee3a4d83b6032436ad61bfa8", + "IPY_MODEL_fc89e3f5fe92417a8fa680d1c679aada", + "IPY_MODEL_8393e3a15697431ca4df2c883a244c47" + ], + "layout": "IPY_MODEL_655d0a5f685c4ba389f7df0aeec1396f" + } + }, + "a1533872921e4eeca8b76a81c1fa8dae": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "a1ed2bfa9147481fa378df313f327426": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_112fe371d54a443faa948d231ba3b778", + "max": 1, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_c0ccdae3037b4aefbda9ae1c44f1798f", + "value": 1 + } + }, + "a7706cbc52be4776bca2511b1f3ceed3": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "a7e1fcb6e01e453a8ae7f56f6faee951": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_57c8abc07c384b9da79a10fccd29e704", + "placeholder": "​", + "style": "IPY_MODEL_96c6e994c9764bbbb158dd945e0f6596", + "value": "merges.txt: " + } + }, + "a821eeb5758d472a9432b154dc4fdd0a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "b005cb13f6e14f98a6cdff7bc7406d6b": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_89a7db6c7b114a299a9a2ada9d9fc633", + "IPY_MODEL_56a942afc0e84db0b9e65cb213a2270a", + "IPY_MODEL_ea1ed850ad4d42949fdfed254e9e2ed1" + ], + "layout": "IPY_MODEL_50dd82810b07495bb85dc2bd82a2d9a9" + } + }, + "b2071d06de0840758aca8184a50461cf": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "b8260fc137e64f15a34a81f6dbacd2f3": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_4f6b9c2772654c29996cf6f98f9b3376", + "placeholder": "​", + "style": "IPY_MODEL_4963c7b82b11449e93df961873181e65", + "value": "added_tokens.json: 100%" + } + }, + "b8c37b59d96542e8b2bd0d1818b67d7c": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "b96eb0324c064ddb99fd611fd16bd7c8": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "b9c63458ed16499dafd94406d6cd9961": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "ba196ff67b0140918c62b20bc0a51cfc": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "bd9a43c626dd4e7a9dff178911772961": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": "20px" + } + }, + "be94e8914efc4aa1887940fa075ff73a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_b96eb0324c064ddb99fd611fd16bd7c8", + "placeholder": "​", + "style": "IPY_MODEL_97b37500e2c549a0990db3ab687a203f", + "value": "vocab.json: " + } + }, + "bf55337bded943a28115b236b5843a58": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "bf7c1535d4e8478593b96e6ffa24f4db": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_b2071d06de0840758aca8184a50461cf", + "placeholder": "​", + "style": "IPY_MODEL_53601dd3a35c4ecea9e0d78881c1163b", + "value": " 11.4M/11.4M [00:00<00:00, 57.1MB/s]" + } + }, + "c02a3f64ae464fddb84805823724447d": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "c0ccdae3037b4aefbda9ae1c44f1798f": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "c37c1b43b04249bfb44ab5794f28c50f": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "c3f10e9e96224ff7ad9b9bd689dcf93e": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "c7c8bc445e7d45f584754f80a5f7a91e": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "caa9a03250c14102af84dcbed12b61f6": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "ccbf6ed7c31b4a11bc31b4020ff7c0cf": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "d4ed60dd9d0e4988986b49fc7c2369ac": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "d5f3e631722641ed87ae7bba300acaaf": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_be94e8914efc4aa1887940fa075ff73a", + "IPY_MODEL_9b4f23f724dd4d2fbe4e0bced4eed79d", + "IPY_MODEL_f99bf961b8814ab78b01aec776bcd15e" + ], + "layout": "IPY_MODEL_93a7ad1ef9c44f26ba71bcce0a52fe2c" + } + }, + "d6b1d2626c7c4c4097e363669c1e3bc7": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_77a9a3a3564b4a42b516ef73f6b3eff2", + "placeholder": "​", + "style": "IPY_MODEL_49b8c93d9ff44b75b4dffc97bd99cfd7", + "value": "generation_config.json: 100%" + } + }, + "d855da6d39b14ebe93134ef1f89c51d1": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "dce034aa439b4326a49bb81bd5f0f3f9": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "ProgressStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "de05b1ffd7fa42bda542ce3b5a662255": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_e578b95b72de4b62b5b0ac0e5a81c94a", + "IPY_MODEL_a1ed2bfa9147481fa378df313f327426", + "IPY_MODEL_6c0ba6aa76bd491896f0f7572bc6c5c8" + ], + "layout": "IPY_MODEL_e55e5e2f79d14fe48fef5e8f22eb89cf" + } + }, + "df43d85bd2fd46dc82f746409a2a3d32": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_a1533872921e4eeca8b76a81c1fa8dae", + "max": 171, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_3d3575a4cd194727962f09f7d5b0c23d", + "value": 171 + } + }, + "e331b7c239884d0d9fb476146a1ed616": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "e501375a0c4b4da09a9d686147aed523": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "e55e5e2f79d14fe48fef5e8f22eb89cf": { + "model_module": "@jupyter-widgets/base", + "model_module_version": "1.2.0", + "model_name": "LayoutModel", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "e578b95b72de4b62b5b0ac0e5a81c94a": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_893bbc7b1a6a489c857f658439055b02", + "placeholder": "​", + "style": "IPY_MODEL_a821eeb5758d472a9432b154dc4fdd0a", + "value": "tokenizer_config.json: " + } + }, + "e70ce7b46e22482aa3e793b3a5169aae": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_a7e1fcb6e01e453a8ae7f56f6faee951", + "IPY_MODEL_74cba38a20344b80a6239fca5734a418", + "IPY_MODEL_307d7f48695b417a814bf7ade5dccf80" + ], + "layout": "IPY_MODEL_e501375a0c4b4da09a9d686147aed523" + } + }, + "e94c42c578264e8894d99e1dd84249f3": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_a7706cbc52be4776bca2511b1f3ceed3", + "placeholder": "​", + "style": "IPY_MODEL_36ac26d37a904b44a01285f37e0fc685", + "value": " 605/605 [00:00<00:00, 73.0kB/s]" + } + }, + "ea1ed850ad4d42949fdfed254e9e2ed1": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_bf55337bded943a28115b236b5843a58", + "placeholder": "​", + "style": "IPY_MODEL_7a3d6b77271a4eebb277c3c25a39cb41", + "value": " 617/617 [00:00<00:00, 72.1kB/s]" + } + }, + "f637e4e79fa2493c9489eecd965b8ce1": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_37c4302d7ba04d7ebf33222f30dc18b5", + "max": 605, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_ba196ff67b0140918c62b20bc0a51cfc", + "value": 605 + } + }, + "f7aeef0949674a3fb0e7dcf70f07f8e4": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "DescriptionStyleModel", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "f99bf961b8814ab78b01aec776bcd15e": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HTMLModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_3a2c9d65c7f840aba8064f6358468edf", + "placeholder": "​", + "style": "IPY_MODEL_e331b7c239884d0d9fb476146a1ed616", + "value": " 2.78M/? [00:00<00:00, 5.21MB/s]" + } + }, + "fc0db98200be480cb63f9f4447433ee0": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "HBoxModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_d6b1d2626c7c4c4097e363669c1e3bc7", + "IPY_MODEL_df43d85bd2fd46dc82f746409a2a3d32", + "IPY_MODEL_a0a15427a41c440fbcf2abcb6891aa23" + ], + "layout": "IPY_MODEL_8ac96d00937449da9846ea28c77addba" + } + }, + "fc89e3f5fe92417a8fa680d1c679aada": { + "model_module": "@jupyter-widgets/controls", + "model_module_version": "1.5.0", + "model_name": "FloatProgressModel", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_24d4b985e965409e9a01f7eb3ba3f23d", + "max": 3087467144, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_696f3bb68ab643a2b3afdc87a12debb8", + "value": 3087467144 + } + } + } + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/DockerFile b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/DockerFile new file mode 100644 index 0000000000000000000000000000000000000000..6655b5ec1baed49b2ab77d03a237233772ddb8ae --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/DockerFile @@ -0,0 +1,34 @@ +# Use a lightweight Python 3.10 image as the base +FROM python:3.10-slim + +# Hugging Face Spaces requires a non-root user for security. +# We create a user named 'user' with user ID 1000. +RUN useradd -m -u 1000 user + +# Set environment variables to ensure Python output is logged immediately +# and to add the local bin directory to the PATH for pip installations. +ENV PATH="/home/user/.local/bin:$PATH" \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 + +# Switch away from root to the new user +USER user + +# Set the working directory inside the container +WORKDIR /home/user/app + +# Copy the requirements file first to leverage Docker cache layers +COPY --chown=user requirements.txt . + +# Upgrade pip and install the dependencies defined by Role 2 +RUN pip install --no-cache-dir --upgrade pip && \ + pip install --no-cache-dir -r requirements.txt + +# Copy the rest of the application files (app.py, openenv.yaml, engine.py, etc.) +COPY --chown=user . . + +# Expose port 7860, which is the default port Hugging Face routes traffic to +EXPOSE 7860 + +# Start the OpenEnv server using the entrypoint defined in app.py +CMD ["python", "app.py"] diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md index 991a6a7bfcba7747510a50df55f2a62e2281cb1c..55dcc433c9908155a8f5fb93a51169b100921b85 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md @@ -1 +1,9 @@ -# MetaHackathon2026Finals \ No newline at end of file +--- +title: Compiler Brain OpenEnv +emoji: 🧠 +colorFrom: indigo +colorTo: purple +sdk: docker +app_file: app.py +pinned: false +--- diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..4ca6c261121de442a517f2fc7ae5ce02fb6b7ac6 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt @@ -0,0 +1,24 @@ +# OpenEnv SDK (Mandatory for the hackathon judging environment) +openenv + +# Core RL Environment +gymnasium + +# Hugging Face Training Stack (Versions matched to your D3 notebook) +trl==0.23.1 +transformers==4.57.1 +peft +accelerate +bitsandbytes + +# Unsloth for efficient 4-bit QLoRA training on T4 GPUs +# Note: Unsloth often prefers being installed via their specific pip wheel or git, +# but this is standard for a requirements file. +unsloth + +# PyTorch (Will default to standard compatible version if no index is specified) +torch + +# Logging & Visualizations (For Day 2 judging criteria) +wandb +matplotlib diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/app.py b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/app.py new file mode 100644 index 0000000000000000000000000000000000000000..e8c748e0423b20310f1295fc7c849bd997dd807e --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/app.py @@ -0,0 +1,182 @@ +import json +import logging +import traceback +from typing import Any, Dict, List + +# OpenEnv SDK import +from openenv import MCPEnvironment + +# ============================================================================== +# ROLE 1 & 3 IMPORTS +# TODO: Import the actual execution engine and generator from your teammates +# ============================================================================== +# from engine import execute_tac, verify_equivalence +# from curriculum import generate_level_code + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("CompilerEnvServer") + +class CompilerEnv(MCPEnvironment): + """ + The OpenEnv Server Wrapper for the Toy-IR Compiler Pass Optimizer. + Acts as the referee between the LLM client and Role 1's Execution Engine. + """ + + def __init__(self): + super().__init__() + # State variables + self.raw_json_code = None + self.current_state_string = "" + self.initial_cycles = 0 + self.current_step = 0 + self.max_steps = 10 # Set a max step limit per episode + + # Cycle Weights (The Physics defined by Role 2) + self.cycle_weights = { + "ADD": 1, + "SUB": 1, + "MUL": 4, + "DIV": 10, + "MEM_LOAD": 20, + "STORE": 20 + } + + def reset(self) -> str: + """ + Grabs unoptimized code, calculates baseline cycles, and translates the + state for the LLM. + """ + self.current_step = 0 + + # 1. Grab new unoptimized code (Role 3 integration) + # TODO: Replace with real generator: self.raw_json_code = generate_level_code() + self.raw_json_code = self._mock_generator() + + # 2. Get baseline cycles (Role 1 integration) + # TODO: Replace with real engine: self.initial_cycles, _ = execute_tac(self.raw_json_code, []) + self.initial_cycles = 100 + + # 3. Translate to Pseudo-Assembly to prevent Attention Dilution + self.current_state_string = self._translate_state(self.raw_json_code) + + logger.info(f"Environment Reset. Baseline Cycles: {self.initial_cycles}") + return self.state() + + def step(self, action: str) -> Dict[str, Any]: + """ + Executes the LLM's chosen optimization pass, verifies math equivalence, + and calculates the reward. + """ + self.current_step += 1 + + # 1. Parse LLM Action (Regex/JSON robustness) + try: + # Assuming the LLM outputs a single pass name or a list of passes + action_data = json.loads(action) + if isinstance(action_data, str): + action_array = [action_data] + else: + action_array = action_data + format_bonus = 0.1 + except json.JSONDecodeError: + # Format Trap + return self._build_step_response( + reward=-2.5, + done=True, + error="Invalid JSON. You must output a valid JSON array of strings." + ) + + # 2. Execute Code & Verify Equivalence (Role 1 Integration) + # TODO: new_cycles, optimized_code = execute_tac(self.raw_json_code, action_array) + # TODO: is_valid = verify_equivalence(self.raw_json_code, optimized_code) + new_cycles = 80 # Mock Data + is_valid = True # Mock Data + + # 3. Calculate Reward Physics + if not is_valid: + # Correctness Penalty + Micro-Variance for GRPO + penalty = -2.0 - (len(action_array) * 0.01) + return self._build_step_response( + reward=penalty, + done=True, + error=f"Code equivalence broken by passes: {action_array}" + ) + + # Calculate Improvement Ratio + Time Tax (-1.0) + cycle_improvement_ratio = (self.initial_cycles - new_cycles) / self.initial_cycles + time_tax = -0.05 * self.current_step # Small tax to prevent pass spamming + reward = cycle_improvement_ratio + format_bonus + time_tax + + # Update state if sequential, or finish if one-shot + # NOTE: For hackathon speed, we treat this as a One-Shot episode + done = True + + return self._build_step_response( + reward=reward, + done=done, + info={"status": "success", "optimized_cycles": new_cycles} + ) + + def state(self) -> str: + """ + Returns the current observation to the LLM. + """ + return f"Current Step: {self.current_step}/{self.max_steps}\n\n{self.current_state_string}" + + def _translate_state(self, raw_json: List[Dict]) -> str: + """ + Translates raw AST JSON into clean pseudo-assembly. + Strips all UUIDs and AST metadata. + """ + pseudo_assembly = [] + instruction_count = 1 + + for inst in raw_json: + op = inst.get("op", "UNKNOWN") + src1 = inst.get("src1", "") + src2 = inst.get("src2", "") + dest = inst.get("dest", "") + + # Format arguments cleanly + args = f"{src1}" if src2 is None else f"{src1}, {src2}" + + if dest: + line = f"{instruction_count}. {dest} = {op} {args}" + else: + line = f"{instruction_count}. {op} {args}" + + pseudo_assembly.append(line) + instruction_count += 1 + + return "\n".join(pseudo_assembly) + + def _build_step_response(self, reward: float, done: bool, error: str = None, info: dict = None) -> Dict[str, Any]: + """Helper to format the standard OpenEnv step return dictionary.""" + response = { + "reward": reward, + "done": done, + "state": self.state() + } + if error: + response["error"] = error + if info: + response["info"] = info + return response + + def _mock_generator(self): + """Mock data so the server runs before Role 1 integrates their engine.""" + return [ + {"op": "CONST", "dest": "a", "src1": 2, "src2": None}, + {"op": "CONST", "dest": "b", "src1": 3, "src2": None}, + {"op": "ADD", "dest": "c", "src1": "a", "src2": "b"} + ] + +if __name__ == "__main__": + logger.info("Initializing CompilerEnv Server...") + try: + env = CompilerEnv() + # openenv.run() or start() depending on the specific MCP wrapper version + env.start() + except Exception as e: + logger.error(f"Failed to start environment server: {e}") + logger.error(traceback.format_exc()) diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_tetris (1).ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_tetris (1).ipynb new file mode 100644 index 0000000000000000000000000000000000000000..c061d098fcd03d7a5ed80d42de43110570714b51 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_tetris (1).ipynb @@ -0,0 +1,651 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "b7fdb9ab", + "metadata": {}, + "outputs": [], + "source": [ + "# ==== CELL 1: Setup ====\n", + "!pip install -q pytest\n", + "import os, sys\n", + "\n", + "LOCAL_DIR = \"/content/toyir\"\n", + "DRIVE_DIR = \"/content/drive/MyDrive/env/toyir\"\n", + "\n", + "os.makedirs(LOCAL_DIR, exist_ok=True)\n", + "sys.path.insert(0, LOCAL_DIR)\n", + "print(\"Workspace ready at\", LOCAL_DIR)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d0337367", + "metadata": {}, + "outputs": [], + "source": [ + "# ==== CELL 1b: Mount Drive + sync workspace ====\n", + "from google.colab import drive\n", + "drive.mount('/content/drive')\n", + "\n", + "import shutil\n", + "os.makedirs(DRIVE_DIR, exist_ok=True)\n", + "\n", + "# On fresh runtime: pull existing files from Drive -> local\n", + "for f in os.listdir(DRIVE_DIR):\n", + " src = os.path.join(DRIVE_DIR, f)\n", + " if os.path.isfile(src) and f.endswith(\".py\"):\n", + " shutil.copy2(src, os.path.join(LOCAL_DIR, f))\n", + "print(\"Synced from Drive:\", sorted(os.listdir(LOCAL_DIR)))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2c83c529", + "metadata": {}, + "outputs": [], + "source": [ + "%%writefile /content/toyir/toy_vm.py\n", + "\"\"\"Toy-IR VM: deterministic, dependency-light execution engine for RLVR.\"\"\"\n", + "from __future__ import annotations\n", + "from copy import deepcopy\n", + "from typing import Any\n", + "\n", + "CYCLE_COST: dict[str, int] = {\n", + " \"CONST\": 1, \"ADD\": 1, \"SUB\": 1, \"MUL\": 3, \"DIV\": 5,\n", + " \"COPY\": 1, \"LOAD\": 4, \"STORE\": 4, \"NOP\": 0,\n", + "}\n", + "OP_ORDER = (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"COPY\", \"LOAD\", \"STORE\", \"NOP\")\n", + "VALID_OPS = frozenset(CYCLE_COST)\n", + "REQUIRED_KEYS = (\"op\", \"dest\", \"src1\", \"src2\")\n", + "\n", + "\n", + "def _is_str_or_int_or_none(x: Any) -> bool:\n", + " return x is None or isinstance(x, str) or (isinstance(x, int) and not isinstance(x, bool))\n", + "\n", + "\n", + "def _is_str_or_none(x: Any) -> bool:\n", + " return x is None or isinstance(x, str)\n", + "\n", + "\n", + "def validate_ir(program: list[dict]) -> dict:\n", + " errors: list[str] = []\n", + " for i, ins in enumerate(program):\n", + " if not isinstance(ins, dict):\n", + " errors.append(f\"[{i}] not a dict\")\n", + " continue\n", + " for k in REQUIRED_KEYS:\n", + " if k not in ins:\n", + " errors.append(f\"[{i}] missing key '{k}'\")\n", + " if \"op\" not in ins:\n", + " continue\n", + " op = ins[\"op\"]\n", + " if op not in VALID_OPS:\n", + " errors.append(f\"[{i}] invalid op '{op}'\")\n", + " continue\n", + " if not _is_str_or_none(ins.get(\"dest\")):\n", + " errors.append(f\"[{i}] dest must be str|None\")\n", + " if not _is_str_or_int_or_none(ins.get(\"src1\")):\n", + " errors.append(f\"[{i}] src1 must be str|int|None\")\n", + " if not _is_str_or_int_or_none(ins.get(\"src2\")):\n", + " errors.append(f\"[{i}] src2 must be str|int|None\")\n", + "\n", + " # per-op contracts\n", + " if op in (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"COPY\", \"LOAD\"):\n", + " if not isinstance(ins.get(\"dest\"), str):\n", + " errors.append(f\"[{i}] {op} requires str dest\")\n", + " if op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n", + " if ins.get(\"src1\") is None or ins.get(\"src2\") is None:\n", + " errors.append(f\"[{i}] {op} requires src1 and src2\")\n", + " if op in (\"CONST\", \"COPY\"):\n", + " if ins.get(\"src1\") is None:\n", + " errors.append(f\"[{i}] {op} requires src1\")\n", + " if op in (\"CONST\", \"COPY\", \"LOAD\", \"STORE\"):\n", + " if ins.get(\"src2\") is not None:\n", + " errors.append(f\"[{i}] {op} requires src2=None\")\n", + " if op == \"LOAD\":\n", + " if not isinstance(ins.get(\"src1\"), str):\n", + " errors.append(f\"[{i}] LOAD src1 must be str (variable name)\")\n", + " if op == \"STORE\":\n", + " if not isinstance(ins.get(\"dest\"), str):\n", + " errors.append(f\"[{i}] STORE dest must be str (variable name)\")\n", + " if ins.get(\"src1\") is None:\n", + " errors.append(f\"[{i}] STORE requires src1\")\n", + " if op == \"NOP\":\n", + " if any(ins.get(k) is not None for k in (\"dest\", \"src1\", \"src2\")):\n", + " errors.append(f\"[{i}] NOP fields must all be None\")\n", + "\n", + " return {\"valid\": len(errors) == 0, \"errors\": errors}\n", + "\n", + "\n", + "def _resolve(operand: Any, variables: dict) -> int:\n", + " if isinstance(operand, str):\n", + " if operand not in variables:\n", + " raise KeyError(f\"undefined variable '{operand}'\")\n", + " return variables[operand]\n", + " if isinstance(operand, int) and not isinstance(operand, bool):\n", + " return operand\n", + " raise TypeError(f\"bad operand: {operand!r}\")\n", + "\n", + "\n", + "def _exec_one(ins: dict, variables: dict, memory: dict) -> int:\n", + " op = ins[\"op\"]\n", + " dest = ins[\"dest\"]\n", + " s1, s2 = ins[\"src1\"], ins[\"src2\"]\n", + "\n", + " if op == \"CONST\":\n", + " variables[dest] = _resolve(s1, variables)\n", + " elif op == \"ADD\":\n", + " variables[dest] = _resolve(s1, variables) + _resolve(s2, variables)\n", + " elif op == \"SUB\":\n", + " variables[dest] = _resolve(s1, variables) - _resolve(s2, variables)\n", + " elif op == \"MUL\":\n", + " variables[dest] = _resolve(s1, variables) * _resolve(s2, variables)\n", + " elif op == \"DIV\":\n", + " b = _resolve(s2, variables)\n", + " if b == 0:\n", + " raise ZeroDivisionError(\"DIV by zero\")\n", + " variables[dest] = _resolve(s1, variables) // b\n", + " elif op == \"COPY\":\n", + " variables[dest] = _resolve(s1, variables)\n", + " elif op == \"LOAD\":\n", + " if not isinstance(s1, str):\n", + " raise TypeError(\"LOAD src1 must be a variable name\")\n", + " addr = variables.get(s1)\n", + " if not (isinstance(addr, int) and not isinstance(addr, bool)):\n", + " raise TypeError(f\"LOAD address from '{s1}' is not int: {addr!r}\")\n", + " if addr not in memory:\n", + " raise KeyError(f\"LOAD from uninitialized addr {addr}\")\n", + " variables[dest] = memory[addr]\n", + " elif op == \"STORE\":\n", + " if not isinstance(dest, str):\n", + " raise TypeError(\"STORE dest must be a variable name\")\n", + " addr = variables.get(dest)\n", + " if not (isinstance(addr, int) and not isinstance(addr, bool)):\n", + " raise TypeError(f\"STORE address from '{dest}' is not int: {addr!r}\")\n", + " memory[addr] = _resolve(s1, variables)\n", + " elif op == \"NOP\":\n", + " pass\n", + " else:\n", + " raise ValueError(f\"non-executable op {op}\")\n", + " return CYCLE_COST[op]\n", + "\n", + "\n", + "def execute(program: list[dict], initial_state: dict) -> dict:\n", + " variables: dict = {}\n", + " memory: dict = {}\n", + " cycles = 0\n", + " try:\n", + " if not isinstance(initial_state, dict):\n", + " raise TypeError(\"initial_state must be a dict\")\n", + " state = deepcopy(initial_state)\n", + " variables = state.get(\"variables\", {})\n", + " memory = state.get(\"memory\", {})\n", + " if not isinstance(variables, dict):\n", + " raise TypeError(\"initial_state['variables'] must be a dict\")\n", + " if not isinstance(memory, dict):\n", + " raise TypeError(\"initial_state['memory'] must be a dict\")\n", + "\n", + " for ins in program:\n", + " cycles += _exec_one(ins, variables, memory)\n", + " except Exception as e:\n", + " return {\"memory\": memory, \"variables\": variables,\n", + " \"cycles\": None, \"success\": False, \"error\": str(e)}\n", + " return {\"memory\": memory, \"variables\": variables,\n", + " \"cycles\": cycles, \"success\": True, \"error\": None}\n", + "\n", + "\n", + "def count_cycles(program: list[dict]) -> int:\n", + " total = 0\n", + " for i, ins in enumerate(program):\n", + " op = ins.get(\"op\") if isinstance(ins, dict) else None\n", + " if op not in CYCLE_COST:\n", + " raise ValueError(f\"count_cycles: invalid or missing op at index {i}: {op!r}\")\n", + " total += CYCLE_COST[op]\n", + " return total\n", + "\n", + "\n", + "def profile(program: list[dict]) -> dict:\n", + " out: dict[str, int] = {\"n_instructions\": len(program)}\n", + " for op in OP_ORDER:\n", + " out[f\"n_{op.lower()}\"] = 0\n", + " for ins in program:\n", + " key = f\"n_{ins['op'].lower()}\"\n", + " if key in out:\n", + " out[key] += 1\n", + " return out\n", + "\n", + "\n", + "def _fmt_operand(x: Any) -> str:\n", + " if x is None:\n", + " return \"\"\n", + " return str(x)\n", + "\n", + "\n", + "def dump_ir(program: list[dict]) -> str:\n", + " stores = [ins[\"dest\"] for ins in program if ins[\"op\"] == \"STORE\"]\n", + " seen, ordered = set(), []\n", + " for s in stores:\n", + " if s not in seen:\n", + " seen.add(s); ordered.append(s)\n", + " header = \"// OBSERVABLE OUT: \" + (\", \".join(f\"mem[{s}]\" for s in ordered) if ordered else \"(none)\")\n", + " lines = [header]\n", + " for ins in program:\n", + " op = ins[\"op\"]\n", + " dest, s1, s2 = ins[\"dest\"], ins[\"src1\"], ins[\"src2\"]\n", + " if op == \"STORE\":\n", + " lines.append(f\"STORE [{dest}] {_fmt_operand(s1)}\")\n", + " elif op == \"LOAD\":\n", + " lines.append(f\"{dest} = LOAD [{_fmt_operand(s1)}]\")\n", + " elif op == \"NOP\":\n", + " lines.append(\"NOP\")\n", + " else:\n", + " parts = [op, _fmt_operand(s1)]\n", + " if s2 is not None:\n", + " parts.append(_fmt_operand(s2))\n", + " lines.append(f\"{dest} = {' '.join(p for p in parts if p)}\")\n", + " return \"\\n\".join(lines)\n", + "\n", + "\n", + "def clean_nops(program: list[dict]) -> list[dict]:\n", + " return [deepcopy(ins) for ins in program if ins[\"op\"] != \"NOP\"]\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3c85a58b", + "metadata": {}, + "outputs": [], + "source": [ + "%%writefile /content/toyir/verifier.py\n", + "\"\"\"Strict equivalence verifier - anti-cheat for RL reward hacking.\"\"\"\n", + "from __future__ import annotations\n", + "from toy_vm import execute\n", + "\n", + "\n", + "def verify(original_program: list[dict],\n", + " optimized_program: list[dict],\n", + " initial_states: list[dict]) -> dict:\n", + " if not initial_states:\n", + " raise ValueError(\"verify requires at least one initial_state\")\n", + "\n", + " results = []\n", + " all_match = True\n", + " for st in initial_states:\n", + " r_o = execute(original_program, st)\n", + " r_p = execute(optimized_program, st)\n", + " mem_o = r_o[\"memory\"] if r_o[\"success\"] else None\n", + " mem_p = r_p[\"memory\"] if r_p[\"success\"] else None\n", + "\n", + " execution_errors = {}\n", + " if not r_o[\"success\"]:\n", + " execution_errors[\"original_error\"] = r_o[\"error\"]\n", + " if not r_p[\"success\"]:\n", + " execution_errors[\"optimized_error\"] = r_p[\"error\"]\n", + "\n", + " if execution_errors:\n", + " match = False\n", + " mismatches = [\"__execution_error__\"]\n", + " else:\n", + " keys = set(mem_o.keys()) | set(mem_p.keys())\n", + " mismatches = [k for k in keys if mem_o.get(k) != mem_p.get(k)]\n", + " match = len(mismatches) == 0\n", + "\n", + " if not match:\n", + " all_match = False\n", + " results.append({\n", + " \"initial_state\": st,\n", + " \"original_memory\": mem_o,\n", + " \"optimized_memory\": mem_p,\n", + " \"match\": match,\n", + " \"mismatches\": sorted(mismatches, key=lambda x: (isinstance(x, str), x)),\n", + " \"execution_errors\": execution_errors or None,\n", + " })\n", + " return {\"equivalent\": all_match, \"results\": results}\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cd7f8fa2", + "metadata": {}, + "outputs": [], + "source": [ + "%%writefile /content/toyir/reward_utils.py\n", + "\"\"\"Reward signal: cycle-reduction metric for RL agent.\"\"\"\n", + "from __future__ import annotations\n", + "from toy_vm import execute\n", + "\n", + "\n", + "def compute_reduction(original_program: list[dict],\n", + " optimized_program: list[dict],\n", + " initial_state: dict) -> dict:\n", + " r_o = execute(original_program, initial_state)\n", + " r_p = execute(optimized_program, initial_state)\n", + "\n", + " if not (r_o[\"success\"] and r_p[\"success\"]):\n", + " return {\n", + " \"original_cycles\": r_o[\"cycles\"],\n", + " \"optimized_cycles\": r_p[\"cycles\"],\n", + " \"absolute_savings\": None,\n", + " \"relative_savings\": None,\n", + " \"success\": False,\n", + " \"error\": r_o[\"error\"] or r_p[\"error\"],\n", + " }\n", + "\n", + " oc, pc = r_o[\"cycles\"], r_p[\"cycles\"]\n", + " rel = 0.0 if oc == 0 else (oc - pc) / oc\n", + " rel = max(-1.0, min(1.0, rel))\n", + " return {\n", + " \"original_cycles\": oc,\n", + " \"optimized_cycles\": pc,\n", + " \"absolute_savings\": oc - pc,\n", + " \"relative_savings\": rel,\n", + " \"success\": True,\n", + " \"error\": None,\n", + " }\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f2a941c", + "metadata": {}, + "outputs": [], + "source": [ + "%%writefile /content/toyir/test_vm.py\n", + "\"\"\"Pytest suite covering VM, verifier, reward, and RL exploit guardrails.\"\"\"\n", + "from __future__ import annotations\n", + "import random\n", + "import pytest\n", + "from toy_vm import (\n", + " validate_ir, execute, count_cycles, profile, dump_ir, clean_nops, OP_ORDER,\n", + ")\n", + "from verifier import verify\n", + "from reward_utils import compute_reduction\n", + "\n", + "\n", + "def I(op, dest=None, src1=None, src2=None):\n", + " return {\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2}\n", + "\n", + "\n", + "def test_const_fold_equivalence():\n", + " orig = [I(\"CONST\", \"t0\", 5), I(\"CONST\", \"t1\", 7), I(\"ADD\", \"t2\", \"t0\", \"t1\"),\n", + " I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"t2\")]\n", + " opt = [I(\"CONST\", \"t2\", 12), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"t2\")]\n", + " res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"]\n", + " red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"absolute_savings\"] > 0\n", + "\n", + "\n", + "def test_dead_code_invisible():\n", + " orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 3),\n", + " I(\"MUL\", \"dead\", \"x\", 99), I(\"STORE\", \"addr\", \"x\")]\n", + " opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 3), I(\"STORE\", \"addr\", \"x\")]\n", + " res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"]\n", + "\n", + "\n", + "def test_fold_then_dce():\n", + " orig = [I(\"CONST\", \"a\", 2), I(\"CONST\", \"b\", 3), I(\"ADD\", \"c\", \"a\", \"b\"),\n", + " I(\"MUL\", \"d\", \"c\", 10), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"d\")]\n", + " opt = [I(\"CONST\", \"d\", 50), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"d\")]\n", + " res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"]\n", + " red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"original_cycles\"] > red[\"optimized_cycles\"]\n", + "\n", + "\n", + "def test_strength_reduction():\n", + " orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 7),\n", + " I(\"MUL\", \"y\", \"x\", 2), I(\"STORE\", \"addr\", \"y\")]\n", + " opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 7),\n", + " I(\"ADD\", \"y\", \"x\", \"x\"), I(\"STORE\", \"addr\", \"y\")]\n", + " res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"]\n", + " red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"absolute_savings\"] > 0\n", + "\n", + "\n", + "def test_already_optimal():\n", + " p = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n", + " red = compute_reduction(p, p, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"absolute_savings\"] == 0\n", + " assert red[\"relative_savings\"] == 0.0\n", + "\n", + "\n", + "def test_div_by_zero():\n", + " p = [I(\"CONST\", \"a\", 10), I(\"CONST\", \"b\", 0), I(\"DIV\", \"c\", \"a\", \"b\")]\n", + " r = execute(p, {\"variables\": {}, \"memory\": {}})\n", + " assert r[\"success\"] is False\n", + " assert \"zero\" in r[\"error\"].lower()\n", + "\n", + "\n", + "def test_hardcoded_exploit_caught():\n", + " orig = [I(\"CONST\", \"addr\", 0), I(\"LOAD\", \"x\", \"in_addr\"),\n", + " I(\"MUL\", \"y\", \"x\", 3), I(\"STORE\", \"addr\", \"y\")]\n", + " opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"y\", 42), I(\"STORE\", \"addr\", \"y\")]\n", + " rng = random.Random(0)\n", + " states = [{\"variables\": {\"in_addr\": 1}, \"memory\": {1: rng.randint(0, 100)}} for _ in range(5)]\n", + " res = verify(orig, opt, states)\n", + " assert res[\"equivalent\"] is False\n", + " assert any(not r[\"match\"] for r in res[\"results\"])\n", + "\n", + "\n", + "def test_nop_zero_cost():\n", + " p = [I(\"NOP\"), I(\"CONST\", \"x\", 1), I(\"NOP\")]\n", + " assert count_cycles(p) == 1\n", + " r = execute(p, {\"variables\": {}, \"memory\": {}})\n", + " assert r[\"success\"] and r[\"cycles\"] == 1\n", + "\n", + "\n", + "def test_load_store_non_int_address():\n", + " p = [I(\"COPY\", \"addr\", \"bad\"), I(\"LOAD\", \"x\", \"addr\")]\n", + " r = execute(p, {\"variables\": {\"bad\": \"notint\"}, \"memory\": {}})\n", + " assert r[\"success\"] is False\n", + "\n", + "\n", + "def test_validate_ir_malformed():\n", + " assert not validate_ir([I(\"FOOBAR\", \"t\", 1, 2)])[\"valid\"]\n", + " assert not validate_ir([{\"op\": \"ADD\", \"dest\": \"x\"}])[\"valid\"]\n", + " assert not validate_ir([I(\"CONST\", 123, 1)])[\"valid\"]\n", + " assert not validate_ir([I(\"STORE\", None, \"x\")])[\"valid\"]\n", + " assert not validate_ir([I(\"LOAD\", \"x\", 42)])[\"valid\"]\n", + " assert not validate_ir([I(\"NOP\", \"x\", 1)])[\"valid\"]\n", + " assert not validate_ir([I(\"ADD\", \"x\", \"a\", None)])[\"valid\"]\n", + " assert not validate_ir([I(\"LOAD\", \"x\", \"addr\", 9)])[\"valid\"]\n", + " assert not validate_ir([I(\"STORE\", \"addr\", \"x\", 9)])[\"valid\"]\n", + " assert validate_ir([I(\"CONST\", \"x\", 1)])[\"valid\"]\n", + "\n", + "\n", + "def test_dump_ir_format():\n", + " p = [I(\"CONST\", \"t0\", 5), I(\"ADD\", \"t1\", \"t0\", \"a\"),\n", + " I(\"MUL\", \"t2\", \"t1\", 3), I(\"STORE\", \"addr0\", \"t2\")]\n", + " s = dump_ir(p)\n", + " lines = s.split(\"\\n\")\n", + " assert lines[0] == \"// OBSERVABLE OUT: mem[addr0]\"\n", + " assert lines[1] == \"t0 = CONST 5\"\n", + " assert lines[2] == \"t1 = ADD t0 a\"\n", + " assert lines[3] == \"t2 = MUL t1 3\"\n", + " assert lines[4] == \"STORE [addr0] t2\"\n", + " assert len(lines) == 5\n", + "\n", + "\n", + "def test_dump_ir_multi_store_header_order():\n", + " p = [I(\"CONST\", \"a\", 0), I(\"CONST\", \"b\", 1),\n", + " I(\"STORE\", \"a\", 5), I(\"STORE\", \"b\", 6), I(\"STORE\", \"a\", 7)]\n", + " s = dump_ir(p)\n", + " assert s.split(\"\\n\")[0] == \"// OBSERVABLE OUT: mem[a], mem[b]\"\n", + "\n", + "\n", + "def test_clean_nops():\n", + " p = [I(\"NOP\"), I(\"CONST\", \"x\", 1), I(\"NOP\"), I(\"CONST\", \"y\", 2), I(\"NOP\")]\n", + " out = clean_nops(p)\n", + " assert len(out) == 2\n", + " assert all(ins[\"op\"] != \"NOP\" for ins in out)\n", + "\n", + "\n", + "def test_verify_mismatches():\n", + " orig = [I(\"CONST\", \"a0\", 0), I(\"CONST\", \"a1\", 1),\n", + " I(\"CONST\", \"x\", 5), I(\"CONST\", \"y\", 9),\n", + " I(\"STORE\", \"a0\", \"x\"), I(\"STORE\", \"a1\", \"y\")]\n", + " opt = [I(\"CONST\", \"a0\", 0), I(\"CONST\", \"a1\", 1),\n", + " I(\"CONST\", \"x\", 5), I(\"CONST\", \"y\", 8),\n", + " I(\"STORE\", \"a0\", \"x\"), I(\"STORE\", \"a1\", \"y\")]\n", + " res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"] is False\n", + " assert res[\"results\"][0][\"mismatches\"] == [1]\n", + "\n", + "\n", + "def test_profile_counts_and_ordering():\n", + " p = [I(\"CONST\", \"x\", 1), I(\"ADD\", \"y\", \"x\", 2), I(\"MUL\", \"z\", \"y\", 3)]\n", + " pr = profile(p)\n", + " assert pr[\"n_instructions\"] == 3\n", + " assert pr[\"n_const\"] == 1 and pr[\"n_add\"] == 1 and pr[\"n_mul\"] == 1\n", + " expected = [\"n_instructions\"] + [f\"n_{op.lower()}\" for op in OP_ORDER]\n", + " assert list(pr.keys()) == expected\n", + "\n", + "\n", + "def test_no_input_mutation():\n", + " p = [I(\"CONST\", \"x\", 1), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"x\")]\n", + " st = {\"variables\": {}, \"memory\": {}}\n", + " snap = {\"variables\": dict(st[\"variables\"]), \"memory\": dict(st[\"memory\"])}\n", + " execute(p, st)\n", + " assert st == snap\n", + "\n", + "\n", + "def test_clean_nops_no_aliasing():\n", + " p = [I(\"CONST\", \"x\", 1), I(\"NOP\")]\n", + " out = clean_nops(p)\n", + " out[0][\"op\"] = \"MUL\"\n", + " assert p[0][\"op\"] == \"CONST\"\n", + "\n", + "\n", + "def test_verify_rejects_empty_states():\n", + " with pytest.raises(ValueError):\n", + " verify([I(\"NOP\")], [I(\"NOP\")], [])\n", + "\n", + "\n", + "def test_verify_surfaces_execution_errors():\n", + " good = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n", + " bad = [I(\"LOAD\", \"x\", \"missing_addr\")]\n", + " res = verify(good, bad, [{\"variables\": {}, \"memory\": {}}])\n", + " assert res[\"equivalent\"] is False\n", + " assert res[\"results\"][0][\"mismatches\"] == [\"__execution_error__\"]\n", + " assert res[\"results\"][0][\"execution_errors\"][\"optimized_error\"] is not None\n", + " assert \"original_error\" not in res[\"results\"][0][\"execution_errors\"]\n", + "\n", + "\n", + "def test_execute_rejects_bad_initial_state_shape():\n", + " r = execute([I(\"CONST\", \"x\", 1)], \"not-a-dict\")\n", + " assert r[\"success\"] is False\n", + " assert \"initial_state\" in r[\"error\"]\n", + "\n", + "\n", + "def test_compute_reduction_failure_path():\n", + " orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n", + " bad = [I(\"CONST\", \"a\", 1), I(\"CONST\", \"b\", 0), I(\"DIV\", \"c\", \"a\", \"b\")]\n", + " red = compute_reduction(orig, bad, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"success\"] is False\n", + " assert red[\"relative_savings\"] is None\n", + " assert red[\"absolute_savings\"] is None\n", + " assert \"zero\" in red[\"error\"].lower()\n", + "\n", + "\n", + "def test_relative_savings_clamped_on_slowdown():\n", + " fast = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n", + " # slower variant: extra dead MULs balloon the cycle count\n", + " slow = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1)]\n", + " slow += [I(\"MUL\", \"junk\", \"x\", 2)] * 50\n", + " slow += [I(\"STORE\", \"addr\", \"x\")]\n", + " red = compute_reduction(fast, slow, {\"variables\": {}, \"memory\": {}})\n", + " assert red[\"absolute_savings\"] < 0\n", + " assert red[\"relative_savings\"] == -1.0 # clamped\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6b83243", + "metadata": {}, + "outputs": [], + "source": [ + "# ==== CELL 6: Push to Drive ====\n", + "import shutil, os\n", + "for f in os.listdir(LOCAL_DIR):\n", + " if f.endswith(\".py\"):\n", + " shutil.copy2(os.path.join(LOCAL_DIR, f), os.path.join(DRIVE_DIR, f))\n", + "print(\"Saved to Drive:\", sorted(os.listdir(DRIVE_DIR)))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5675f9a0", + "metadata": {}, + "outputs": [], + "source": [ + "# ==== CELL 7: Run tests ====\n", + "import subprocess\n", + "res = subprocess.run(\n", + " [\"python\", \"-m\", \"pytest\", \"-v\", \"test_vm.py\"],\n", + " cwd=LOCAL_DIR, capture_output=True, text=True,\n", + ")\n", + "print(res.stdout)\n", + "print(res.stderr)\n", + "assert res.returncode == 0, \"Tests failed\"\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a500bb3", + "metadata": {}, + "outputs": [], + "source": [ + "# ==== CELL 8: Smoke demo ====\n", + "import importlib, toy_vm, verifier, reward_utils\n", + "importlib.reload(toy_vm); importlib.reload(verifier); importlib.reload(reward_utils)\n", + "from toy_vm import dump_ir\n", + "from reward_utils import compute_reduction\n", + "from verifier import verify\n", + "\n", + "orig = [\n", + " {\"op\":\"CONST\",\"dest\":\"addr\",\"src1\":0,\"src2\":None},\n", + " {\"op\":\"CONST\",\"dest\":\"a\",\"src1\":2,\"src2\":None},\n", + " {\"op\":\"CONST\",\"dest\":\"b\",\"src1\":3,\"src2\":None},\n", + " {\"op\":\"ADD\",\"dest\":\"c\",\"src1\":\"a\",\"src2\":\"b\"},\n", + " {\"op\":\"MUL\",\"dest\":\"d\",\"src1\":\"c\",\"src2\":2},\n", + " {\"op\":\"MUL\",\"dest\":\"dead\",\"src1\":\"d\",\"src2\":99},\n", + " {\"op\":\"STORE\",\"dest\":\"addr\",\"src1\":\"d\",\"src2\":None},\n", + "]\n", + "opt = [\n", + " {\"op\":\"CONST\",\"dest\":\"addr\",\"src1\":0,\"src2\":None},\n", + " {\"op\":\"CONST\",\"dest\":\"d\",\"src1\":10,\"src2\":None},\n", + " {\"op\":\"STORE\",\"dest\":\"addr\",\"src1\":\"d\",\"src2\":None},\n", + "]\n", + "\n", + "print(dump_ir(orig)); print(\"---\"); print(dump_ir(opt)); print(\"---\")\n", + "print(compute_reduction(orig, opt, {\"variables\":{},\"memory\":{}}))\n", + "print(\"equivalent:\", verify(orig, opt, [{\"variables\":{},\"memory\":{}}])[\"equivalent\"])\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb index 926491c97464967c001cab5b34009ed175a8a465..9db7916ad73b3484de1587691f1ca5aaba53825e 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb @@ -35,12 +35,12 @@ "metadata": { "id": "_XI5jT2Ibvrf" }, - "execution_count": 2, + "execution_count": 30, "outputs": [] }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 31, "metadata": { "id": "j6CQ327KWXwr" }, @@ -163,9 +163,9 @@ "base_uri": "https://localhost:8080/" }, "id": "a86vD1OyWcB9", - "outputId": "2d6d4604-6dc1-4c29-af64-e41cf5717c0d" + "outputId": "2d59b641-be76-4e12-8f7c-06934ab49110" }, - "execution_count": 4, + "execution_count": 32, "outputs": [ { "output_type": "stream", @@ -331,7 +331,7 @@ "metadata": { "id": "CGk8Fz4CZA3t" }, - "execution_count": 5, + "execution_count": 33, "outputs": [] }, { @@ -416,9 +416,9 @@ "base_uri": "https://localhost:8080/" }, "id": "WLKSTVa8fRF8", - "outputId": "e9480e1e-5d59-4afe-a884-bae0fa4a2192" + "outputId": "b039a3da-db4b-4332-9236-8906abc619af" }, - "execution_count": 6, + "execution_count": 34, "outputs": [ { "output_type": "stream", @@ -532,7 +532,7 @@ "metadata": { "id": "5vo-KgaffTFR" }, - "execution_count": 12, + "execution_count": 35, "outputs": [] }, { @@ -618,9 +618,9 @@ "base_uri": "https://localhost:8080/" }, "id": "vv8QRFRlvzdq", - "outputId": "e1d35874-ff1a-4dec-f43f-e22867dd3911" + "outputId": "4f2e9310-7ad9-43ab-cda0-ff1f03e1571b" }, - "execution_count": 8, + "execution_count": 36, "outputs": [ { "output_type": "stream", @@ -758,7 +758,7 @@ "metadata": { "id": "IaczDNeLv5PW" }, - "execution_count": 9, + "execution_count": 37, "outputs": [] }, { @@ -857,9 +857,9 @@ "base_uri": "https://localhost:8080/" }, "id": "DWDUZUQvwXye", - "outputId": "1d1e327f-5f26-43e4-a6ac-f7a286da3781" + "outputId": "f0851e8a-177a-444e-ed1b-a00c6e1e7120" }, - "execution_count": 10, + "execution_count": 38, "outputs": [ { "output_type": "stream", @@ -999,7 +999,7 @@ "metadata": { "id": "yB_yMNOZwaPT" }, - "execution_count": 13, + "execution_count": 39, "outputs": [] }, { @@ -1051,10 +1051,10 @@ "colab": { "base_uri": "https://localhost:8080/" }, - "id": "SUqa7odp3wB_", - "outputId": "c3f44089-213a-4a69-f29f-42b755ed142a" + "id": "nWoFDPTH-Def", + "outputId": "593cef57-8e1b-4e4f-fcf6-13a91c140483" }, - "execution_count": 14, + "execution_count": 40, "outputs": [ { "output_type": "stream", @@ -1115,11 +1115,10 @@ " - 3-5 CONSTs (some literal, some used in arithmetic)\n", " - 3-5 arithmetic ops (ADD, SUB, MUL, DIV) with chaining\n", " - 1-3 dead variables (DCE opportunities)\n", - " - 1 LOAD from initial memory (introduces non-constant variable)\n", + " - 1 LOAD from initial memory whose result is GUARANTEED to flow into\n", + " the final STORE (forces agent to reason around unfoldable runtime values)\n", " - 1 STORE at the end (observable output)\n", - " - Mix of operands designed to hit peephole patterns occasionally:\n", - " - MUL by literal 2 / 1 / 0 (not always β€” diversity matters)\n", - " - DIV by non-zero divisor only (Design Assumption #4)\n", + " - DIV always uses literal non-zero divisor (Design Assumption #4)\n", "\n", " Returns:\n", " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", @@ -1133,7 +1132,7 @@ " var_counter += 1\n", " return name\n", "\n", - " available_vars = [] # vars that exist and can be used as sources\n", + " available_vars = []\n", "\n", " # Step 1: 3-5 CONSTs\n", " num_consts = random.randint(3, 5)\n", @@ -1144,7 +1143,6 @@ " available_vars.append(var)\n", "\n", " # Step 2: 1 LOAD from initial memory\n", - " # We'll seed initial_mem with something at a fresh address.\n", " load_addr_var = \"addr_in\"\n", " loaded_var = new_var()\n", " instructions.append({\n", @@ -1157,22 +1155,16 @@ " last_result = None\n", " for _ in range(num_arith):\n", " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", - "\n", - " # 30% chance to use a literal as src2 to expose peephole opportunities\n", - " # (MUL by 2/1/0, etc.)\n", " use_literal_src2 = random.random() < 0.3\n", "\n", " src1 = random.choice(available_vars)\n", " if use_literal_src2:\n", - " # For DIV, never emit literal 0 (Design Assumption #4)\n", " if op == \"DIV\":\n", - " src2 = random.choice([1, 2, 3, 4]) # safe non-zero divisors\n", + " src2 = random.choice([1, 2, 3, 4])\n", " else:\n", - " src2 = random.choice([0, 1, 2, 3]) # 0/1/2 hit peephole patterns\n", + " src2 = random.choice([0, 1, 2, 3])\n", " else:\n", " src2 = random.choice(available_vars)\n", - " # If we sampled a variable for DIV's src2, we can't be sure it's non-zero\n", - " # at runtime. To stay safe per Assumption #4, restrict DIV to literal src2.\n", " if op == \"DIV\":\n", " src2 = random.choice([1, 2, 3, 4])\n", "\n", @@ -1181,24 +1173,35 @@ " available_vars.append(dest)\n", " last_result = dest\n", "\n", - " # Step 4: 1-3 extra dead CONSTs sprinkled in (DCE targets)\n", + " # Step 4: 1-3 extra dead CONSTs (DCE targets)\n", " num_dead = random.randint(1, 3)\n", " for _ in range(num_dead):\n", " dead_var = new_var()\n", " dead_value = random.randint(1, 20)\n", - " # Insert at a random position before the (eventual) STORE\n", " insert_pos = random.randint(0, len(instructions))\n", " instructions.insert(insert_pos, {\n", " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", " })\n", - " # Note: dead_var intentionally never used afterward\n", "\n", - " # Step 5: STORE the final result\n", + " # === Step 5 (NEW): force LOAD-into-chain dependency ===\n", + " # Inject a final ADD that combines last_result with the loaded variable.\n", + " # This guarantees the LOAD result flows into the STORE β€” DCE can no longer\n", + " # eliminate it, and CF cannot collapse the entire program into a CONST.\n", + " final_dest = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": final_dest,\n", + " \"src1\": last_result,\n", + " \"src2\": loaded_var,\n", + " })\n", + " last_result = final_dest\n", + "\n", + " # Step 6: STORE the final result\n", " initial_vars = {\n", " \"addr0\": 0,\n", - " \"addr_in\": 1, # address that the LOAD reads from\n", + " \"addr_in\": 1,\n", " }\n", - " initial_mem = {1: random.randint(1, 20)} # seed mem[1] with a random value\n", + " initial_mem = {1: random.randint(1, 20)}\n", "\n", " instructions.append({\n", " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", @@ -1212,9 +1215,9 @@ " }" ], "metadata": { - "id": "OL6QbL0H3xuP" + "id": "SUqa7odp3wB_" }, - "execution_count": 15, + "execution_count": 41, "outputs": [] }, { @@ -1255,9 +1258,9 @@ "base_uri": "https://localhost:8080/" }, "id": "D0XjNetU5gIo", - "outputId": "1c35cdfc-a95d-48ff-c7ce-c4c236bbbbdb" + "outputId": "1462cf9a-6f36-49a0-b2e8-bae5e0ef36eb" }, - "execution_count": 16, + "execution_count": 42, "outputs": [ { "output_type": "stream", @@ -1278,8 +1281,9 @@ "9: v6 = v4 + v0 # ADD [1 cycle]\n", "10: v7 = v0 // 2 # DIV [5 cycles]\n", "11: v8 = v0 - v3 # SUB [1 cycle]\n", - "12: mem[addr0] = v8 # STORE [4 cycles]\n", - "// TOTAL: 23 cycles\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", "initial_vars: {'addr0': 0, 'addr_in': 1}\n", "initial_mem : {1: 1}\n", "\n", @@ -1293,8 +1297,9 @@ "5: v4 = v3 // 2 # DIV [5 cycles]\n", "6: v5 = v3 + v3 # ADD [1 cycle]\n", "7: v6 = v2 + v5 # ADD [1 cycle]\n", - "8: mem[addr0] = v6 # STORE [4 cycles]\n", - "// TOTAL: 19 cycles\n", + "8: v8 = v6 + v3 # ADD [1 cycle]\n", + "9: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 20 cycles\n", "initial_vars: {'addr0': 0, 'addr_in': 1}\n", "initial_mem : {1: 11}\n", "\n", @@ -1311,8 +1316,9 @@ "8: v10 = 2 # CONST [1 cycle]\n", "9: v8 = v1 + v3 # ADD [1 cycle]\n", "10: v9 = v6 + v0 # ADD [1 cycle]\n", - "11: mem[addr0] = v9 # STORE [4 cycles]\n", - "// TOTAL: 18 cycles\n", + "11: v11 = v9 + v4 # ADD [1 cycle]\n", + "12: mem[addr0] = v11 # STORE [4 cycles]\n", + "// TOTAL: 19 cycles\n", "initial_vars: {'addr0': 0, 'addr_in': 1}\n", "initial_mem : {1: 5}\n", "\n", @@ -1331,8 +1337,9 @@ "9: v6 = v4 + v0 # ADD [1 cycle]\n", "10: v7 = v0 // 2 # DIV [5 cycles]\n", "11: v8 = v0 - v3 # SUB [1 cycle]\n", - "12: mem[addr0] = v8 # STORE [4 cycles]\n", - "// TOTAL: 23 cycles\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", "\n", "AFTER CF:\n", "// OBSERVABLE OUT: mem[0]\n", @@ -1348,30 +1355,1955 @@ "9: v6 = 6 # CONST [1 cycle]\n", "10: v7 = 1 # CONST [1 cycle]\n", "11: v8 = -2 # CONST [1 cycle]\n", - "12: mem[addr0] = v8 # STORE [4 cycles]\n", - "// TOTAL: 19 cycles\n", + "12: v12 = -2 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 20 cycles\n", "\n", "AFTER CF + DCE:\n", "// OBSERVABLE OUT: mem[0]\n", - "0: v8 = -2 # CONST [1 cycle]\n", - "1: mem[addr0] = v8 # STORE [4 cycles]\n", - "// TOTAL: 5 cycles\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", "\n", "AFTER CF + DCE + PEEPHOLE:\n", "// OBSERVABLE OUT: mem[0]\n", - "0: v8 = -2 # CONST [1 cycle]\n", - "1: mem[addr0] = v8 # STORE [4 cycles]\n", - "// TOTAL: 5 cycles\n" + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n" ] } ] }, { "cell_type": "code", - "source": [], + "source": [ + "# === Compose helper + evaluation harness ===\n", + "\n", + "# Registry mapping pass names (the LLM's action vocabulary) to functions.\n", + "PASS_REGISTRY = {\n", + " \"constant_folding\": constant_folding,\n", + " \"dead_code_elimination\": dead_code_elimination,\n", + " \"peephole_optimization\": peephole_optimization,\n", + "}\n", + "\n", + "\n", + "def count_cycles(instructions):\n", + " \"\"\"Sum cycle costs across an instruction list.\"\"\"\n", + " return sum(CYCLE_COSTS.get(instr[\"op\"], 0) for instr in instructions)\n", + "\n", + "\n", + "def apply_passes(instructions, sequence):\n", + " \"\"\"\n", + " Apply a sequence of passes (by name) to an instruction list.\n", + "\n", + " Args:\n", + " instructions: list of TAC instruction dicts.\n", + " sequence: list of pass-name strings, e.g. [\"constant_folding\", \"dead_code_elimination\"].\n", + "\n", + " Returns:\n", + " new instruction list (does not mutate input).\n", + "\n", + " Raises:\n", + " ValueError if an unknown pass name is provided. (This is intentional β€”\n", + " if Harshal's LLM emits a garbage pass name, we want to know loudly.)\n", + " \"\"\"\n", + " current = [instr.copy() for instr in instructions] # defensive copy\n", + " for pass_name in sequence:\n", + " if pass_name not in PASS_REGISTRY:\n", + " raise ValueError(f\"Unknown pass: {pass_name!r}. Known: {list(PASS_REGISTRY.keys())}\")\n", + " current = PASS_REGISTRY[pass_name](current)\n", + " return current\n", + "\n", + "\n", + "def evaluate(program, sequence):\n", + " \"\"\"\n", + " Run a pass sequence on a program and return cycle counts before/after.\n", + "\n", + " Args:\n", + " program: full program dict (with 'instructions', 'observable_addrs', etc.)\n", + " sequence: list of pass names.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - original_cycles: int\n", + " - optimized_cycles: int\n", + " - optimized_instructions: list[dict]\n", + " - cycle_reduction: int (original - optimized)\n", + " - reduction_pct: float (0.0 to 1.0)\n", + " \"\"\"\n", + " original_instructions = program[\"instructions\"]\n", + " original_cycles = count_cycles(original_instructions)\n", + "\n", + " optimized_instructions = apply_passes(original_instructions, sequence)\n", + " optimized_cycles = count_cycles(optimized_instructions)\n", + "\n", + " cycle_reduction = original_cycles - optimized_cycles\n", + " reduction_pct = cycle_reduction / original_cycles if original_cycles > 0 else 0.0\n", + "\n", + " return {\n", + " \"original_cycles\": original_cycles,\n", + " \"optimized_cycles\": optimized_cycles,\n", + " \"optimized_instructions\": optimized_instructions,\n", + " \"cycle_reduction\": cycle_reduction,\n", + " \"reduction_pct\": reduction_pct,\n", + " }" + ], "metadata": { "id": "Vo-9MitD5hwM" }, + "execution_count": 43, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for compose helper + eval harness ===\n", + "\n", + "# Test 1: count_cycles on a known instruction list\n", + "test1_instrs = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"b\", \"src1\": \"a\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"b\", \"src2\": None},\n", + "]\n", + "print(f\"Test 1 (count_cycles): {count_cycles(test1_instrs)}\")\n", + "# Expected: 1 + 3 + 4 = 8\n", + "\n", + "# Test 2: apply_passes with single pass\n", + "result = apply_passes(test1_instrs, [\"constant_folding\"])\n", + "print(f\"\\nTest 2 (apply CF only):\")\n", + "for i, instr in enumerate(result): print(f\" {i}: {instr}\")\n", + "# Expected: MUL becomes ... actually 5*2 β†’ CONST 10. Wait β€” 'a' is const 5,\n", + "# but src2=2 is literal. Both resolve. CF should fold: b = CONST 10.\n", + "\n", + "# Test 3: apply_passes with sequence\n", + "result = apply_passes(test1_instrs, [\"constant_folding\", \"dead_code_elimination\"])\n", + "print(f\"\\nTest 3 (CF + DCE):\")\n", + "for i, instr in enumerate(result): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 4: full evaluate on a generated Level 2 program\n", + "random.seed(42)\n", + "prog = generate_level_2()\n", + "print(f\"\\nTest 4 (evaluate on Level 2 seed=42):\")\n", + "print(f\"Original program ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "eval_result = evaluate(prog, [\"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\"])\n", + "print(f\"\\nEvaluation result:\")\n", + "print(f\" original_cycles : {eval_result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {eval_result['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {eval_result['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {eval_result['reduction_pct']:.1%}\")\n", + "print(f\"\\nOptimized program:\")\n", + "print(dump_ir({**prog, \"instructions\": eval_result[\"optimized_instructions\"]}))\n", + "\n", + "# Test 5: unknown pass name raises clearly\n", + "try:\n", + " apply_passes(test1_instrs, [\"fake_pass\"])\n", + " print(\"\\nTest 5 (unknown pass): FAIL β€” should have raised\")\n", + "except ValueError as e:\n", + " print(f\"\\nTest 5 (unknown pass raises): PASS β€” {e}\")\n", + "\n", + "# Test 6: empty sequence is a no-op\n", + "result = apply_passes(test1_instrs, [])\n", + "matches = result == test1_instrs\n", + "print(f\"\\nTest 6 (empty sequence is no-op): {'PASS' if matches else 'FAIL'}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "nNLyvg2JA-Xz", + "outputId": "4bb89933-5cd6-4111-9bab-e4d364372972" + }, + "execution_count": 44, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (count_cycles): 8\n", + "\n", + "Test 2 (apply CF only):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 10, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'b', 'src2': None}\n", + "\n", + "Test 3 (CF + DCE):\n", + " 0: {'op': 'CONST', 'dest': 'b', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'b', 'src2': None}\n", + "\n", + "Test 4 (evaluate on Level 2 seed=42):\n", + "Original program (24 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "\n", + "Evaluation result:\n", + " original_cycles : 24\n", + " optimized_cycles : 9\n", + " cycle_reduction : 15\n", + " reduction_pct : 62.5%\n", + "\n", + "Optimized program:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", + "\n", + "Test 5 (unknown pass raises): PASS β€” Unknown pass: 'fake_pass'. Known: ['constant_folding', 'dead_code_elimination', 'peephole_optimization']\n", + "\n", + "Test 6 (empty sequence is no-op): PASS\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Fixed-Point -O3 Baseline ===\n", + "\n", + "# The pass sequence applied in each iteration.\n", + "# Reasoning:\n", + "# CF first : propagates and folds constants, maximizing info for downstream passes\n", + "# PEEP next : rewrites special-value patterns (MUL by 2, ADD with 0) using propagated literals\n", + "# DCE last : sweeps up all newly-dead instructions\n", + "# Then loop until no further changes (fixed point).\n", + "O3_PASS_SEQUENCE = [\n", + " \"constant_folding\",\n", + " \"peephole_optimization\",\n", + " \"dead_code_elimination\",\n", + "]\n", + "\n", + "# Safety cap: prevent infinite loops on pathological inputs.\n", + "# In practice, fixed point is reached in 1-3 iterations.\n", + "MAX_FIXED_POINT_ITERATIONS = 10\n", + "\n", + "\n", + "def apply_o3_baseline(instructions):\n", + " \"\"\"\n", + " Apply the fixed-point -O3 baseline to an instruction list.\n", + "\n", + " Runs O3_PASS_SEQUENCE in a loop until the program stops changing\n", + " (fixed point reached) or MAX_FIXED_POINT_ITERATIONS is hit.\n", + "\n", + " Args:\n", + " instructions: list of TAC instruction dicts.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - optimized_instructions: list[dict] (final program)\n", + " - iterations: int (how many full sequence applications happened)\n", + " - converged: bool (True if fixed point reached, False if iter cap hit)\n", + " \"\"\"\n", + " current = [instr.copy() for instr in instructions]\n", + "\n", + " for iteration in range(1, MAX_FIXED_POINT_ITERATIONS + 1):\n", + " before = current\n", + " current = apply_passes(current, O3_PASS_SEQUENCE)\n", + " if current == before:\n", + " return {\n", + " \"optimized_instructions\": current,\n", + " \"iterations\": iteration,\n", + " \"converged\": True,\n", + " }\n", + "\n", + " return {\n", + " \"optimized_instructions\": current,\n", + " \"iterations\": MAX_FIXED_POINT_ITERATIONS,\n", + " \"converged\": False,\n", + " }\n", + "\n", + "\n", + "def evaluate_baseline(program):\n", + " \"\"\"\n", + " Run the -O3 baseline on a program. Convenience wrapper around apply_o3_baseline\n", + " that also computes cycle metrics.\n", + "\n", + " Args:\n", + " program: full program dict.\n", + "\n", + " Returns:\n", + " dict with original_cycles, optimized_cycles, cycle_reduction, reduction_pct,\n", + " iterations, converged, optimized_instructions.\n", + " \"\"\"\n", + " original_cycles = count_cycles(program[\"instructions\"])\n", + "\n", + " baseline_result = apply_o3_baseline(program[\"instructions\"])\n", + " optimized = baseline_result[\"optimized_instructions\"]\n", + " optimized_cycles = count_cycles(optimized)\n", + "\n", + " return {\n", + " \"original_cycles\": original_cycles,\n", + " \"optimized_cycles\": optimized_cycles,\n", + " \"cycle_reduction\": original_cycles - optimized_cycles,\n", + " \"reduction_pct\": (original_cycles - optimized_cycles) / original_cycles if original_cycles > 0 else 0.0,\n", + " \"iterations\": baseline_result[\"iterations\"],\n", + " \"converged\": baseline_result[\"converged\"],\n", + " \"optimized_instructions\": optimized,\n", + " }" + ], + "metadata": { + "id": "LFQfkN6hBCMS" + }, + "execution_count": 45, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for fixed-point -O3 baseline ===\n", + "\n", + "# Test 1: simple Level 1 program\n", + "random.seed(42)\n", + "prog1 = generate_level_1()\n", + "print(\"Test 1 (-O3 on Level 1, seed=42):\")\n", + "print(f\"Original ({count_cycles(prog1['instructions'])} cycles):\")\n", + "print(dump_ir(prog1))\n", + "\n", + "result1 = evaluate_baseline(prog1)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result1['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result1['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result1['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result1['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result1['iterations']}\")\n", + "print(f\" converged : {result1['converged']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog1, \"instructions\": result1[\"optimized_instructions\"]}))\n", + "\n", + "# Test 2: Level 2 β€” should keep LOAD alive\n", + "print(\"\\n\" + \"=\"*60)\n", + "random.seed(42)\n", + "prog2 = generate_level_2()\n", + "print(\"Test 2 (-O3 on Level 2, seed=42):\")\n", + "print(f\"Original ({count_cycles(prog2['instructions'])} cycles):\")\n", + "print(dump_ir(prog2))\n", + "\n", + "result2 = evaluate_baseline(prog2)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result2['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result2['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result2['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result2['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result2['iterations']}\")\n", + "print(f\" converged : {result2['converged']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog2, \"instructions\": result2[\"optimized_instructions\"]}))\n", + "\n", + "# Test 3: aggregate stats across many programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 2 programs):\")\n", + "total_original = 0\n", + "total_optimized = 0\n", + "iter_distribution = {}\n", + "for seed in range(100, 120):\n", + " random.seed(seed)\n", + " prog = generate_level_2()\n", + " result = evaluate_baseline(prog)\n", + " total_original += result[\"original_cycles\"]\n", + " total_optimized += result[\"optimized_cycles\"]\n", + " iter_distribution[result[\"iterations\"]] = iter_distribution.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_reduction_pct = (total_original - total_optimized) / total_original\n", + "print(f\" total original cycles : {total_original}\")\n", + "print(f\" total optimized cycles : {total_optimized}\")\n", + "print(f\" avg reduction pct : {avg_reduction_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_distribution}\")\n", + "\n", + "# Test 4: idempotence β€” running baseline twice gives same result\n", + "print(\"\\nTest 4 (idempotence β€” applying baseline to its own output is no-op):\")\n", + "random.seed(7)\n", + "prog4 = generate_level_2()\n", + "once = apply_o3_baseline(prog4[\"instructions\"])\n", + "twice = apply_o3_baseline(once[\"optimized_instructions\"])\n", + "matches = once[\"optimized_instructions\"] == twice[\"optimized_instructions\"]\n", + "print(f\" PASS: baseline output is at fixed point\" if matches else \" FAIL: applying baseline twice changed program\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "1cGy1HTSCccB", + "outputId": "61df3175-6aba-4383-e16b-d56e477bf248" + }, + "execution_count": 46, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (-O3 on Level 1, seed=42):\n", + "Original (7 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 1 # CONST [1 cycle]\n", + "1: v1 = 5 # CONST [1 cycle]\n", + "2: v2 = v0 + v0 # ADD [1 cycle]\n", + "3: mem[addr0] = v2 # STORE [4 cycles]\n", + "// TOTAL: 7 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 7\n", + " optimized_cycles : 5\n", + " cycle_reduction : 2\n", + " reduction_pct : 28.6%\n", + " iterations : 2\n", + " converged : True\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v2 = 2 # CONST [1 cycle]\n", + "1: mem[addr0] = v2 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n", + "\n", + "============================================================\n", + "Test 2 (-O3 on Level 2, seed=42):\n", + "Original (24 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: v12 = v8 + v5 # ADD [1 cycle]\n", + "13: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 24 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 24\n", + " optimized_cycles : 9\n", + " cycle_reduction : 15\n", + " reduction_pct : 62.5%\n", + " iterations : 2\n", + " converged : True\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "1: v12 = -2 + v5 # ADD [1 cycle]\n", + "2: mem[addr0] = v12 # STORE [4 cycles]\n", + "// TOTAL: 9 cycles\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 2 programs):\n", + " total original cycles : 522\n", + " total optimized cycles : 231\n", + " avg reduction pct : 55.7%\n", + " iteration distribution : {2: 19, 3: 1}\n", + "\n", + "Test 4 (idempotence β€” applying baseline to its own output is no-op):\n", + " PASS: baseline output is at fixed point\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "# === Held-Out Test Set ===\n", + "\n", + "# Seed partitioning convention (LOCKED β€” must match Harshal's training sampler):\n", + "TRAINING_SEED_MIN = 0\n", + "TRAINING_SEED_MAX = 9999\n", + "HOLDOUT_SEED_MIN = 10000\n", + "HOLDOUT_SEED_MAX = 19999\n", + "\n", + "# Map level number -> generator function.\n", + "# Levels 3-5 will be added as we build them.\n", + "LEVEL_GENERATORS = {\n", + " 1: generate_level_1,\n", + " 2: generate_level_2,\n", + "}\n", + "\n", + "\n", + "def generate_test_set(n_per_level=20, levels=None):\n", + " \"\"\"\n", + " Generate a deterministic held-out test set across difficulty levels.\n", + "\n", + " Uses seeds from the held-out range (10000-19999) only β€” guaranteed\n", + " not to overlap with training data (0-9999).\n", + "\n", + " Args:\n", + " n_per_level: number of programs per difficulty level (default 20).\n", + " levels: list of level numbers to include. Defaults to all available.\n", + "\n", + " Returns:\n", + " list of dicts, each with:\n", + " - level: int (which difficulty level it came from)\n", + " - seed: int (the seed used to generate it)\n", + " - program: full program dict\n", + " \"\"\"\n", + " if levels is None:\n", + " levels = sorted(LEVEL_GENERATORS.keys())\n", + "\n", + " test_set = []\n", + "\n", + " for level in levels:\n", + " if level not in LEVEL_GENERATORS:\n", + " raise ValueError(f\"No generator for level {level}. Available: {list(LEVEL_GENERATORS.keys())}\")\n", + "\n", + " generator = LEVEL_GENERATORS[level]\n", + "\n", + " # Use held-out seeds, offset by level so seeds don't repeat across levels.\n", + " # Level 1 uses 10000-10000+n, Level 2 uses 11000-11000+n, etc.\n", + " base_seed = HOLDOUT_SEED_MIN + (level - 1) * 1000\n", + "\n", + " for i in range(n_per_level):\n", + " seed = base_seed + i\n", + " random.seed(seed)\n", + " program = generator()\n", + " test_set.append({\n", + " \"level\": level,\n", + " \"seed\": seed,\n", + " \"program\": program,\n", + " })\n", + "\n", + " return test_set\n", + "\n", + "\n", + "def run_baseline_on_test_set(test_set):\n", + " \"\"\"\n", + " Run the fixed-point -O3 baseline on every program in the test set.\n", + "\n", + " Args:\n", + " test_set: list of {level, seed, program} dicts.\n", + "\n", + " Returns:\n", + " dict with:\n", + " - per_program: list of per-program results (cycles before/after, level, seed)\n", + " - by_level: dict mapping level -> aggregate stats\n", + " - overall: aggregate stats across all programs\n", + " \"\"\"\n", + " per_program = []\n", + "\n", + " for entry in test_set:\n", + " result = evaluate_baseline(entry[\"program\"])\n", + " per_program.append({\n", + " \"level\": entry[\"level\"],\n", + " \"seed\": entry[\"seed\"],\n", + " \"original_cycles\": result[\"original_cycles\"],\n", + " \"optimized_cycles\": result[\"optimized_cycles\"],\n", + " \"cycle_reduction\": result[\"cycle_reduction\"],\n", + " \"reduction_pct\": result[\"reduction_pct\"],\n", + " \"iterations\": result[\"iterations\"],\n", + " })\n", + "\n", + " # Aggregate by level\n", + " by_level = {}\n", + " for entry in per_program:\n", + " lvl = entry[\"level\"]\n", + " if lvl not in by_level:\n", + " by_level[lvl] = {\n", + " \"n_programs\": 0,\n", + " \"total_original\": 0,\n", + " \"total_optimized\": 0,\n", + " \"sum_reduction_pct\": 0.0,\n", + " }\n", + " by_level[lvl][\"n_programs\"] += 1\n", + " by_level[lvl][\"total_original\"] += entry[\"original_cycles\"]\n", + " by_level[lvl][\"total_optimized\"] += entry[\"optimized_cycles\"]\n", + " by_level[lvl][\"sum_reduction_pct\"] += entry[\"reduction_pct\"]\n", + "\n", + " # Compute averages\n", + " for lvl, stats in by_level.items():\n", + " stats[\"avg_original\"] = stats[\"total_original\"] / stats[\"n_programs\"]\n", + " stats[\"avg_optimized\"] = stats[\"total_optimized\"] / stats[\"n_programs\"]\n", + " stats[\"avg_reduction_pct\"] = stats[\"sum_reduction_pct\"] / stats[\"n_programs\"]\n", + " # Aggregate reduction percent (different from average of percents)\n", + " stats[\"aggregate_reduction_pct\"] = (\n", + " (stats[\"total_original\"] - stats[\"total_optimized\"]) / stats[\"total_original\"]\n", + " if stats[\"total_original\"] > 0 else 0.0\n", + " )\n", + "\n", + " # Overall\n", + " total_orig = sum(e[\"original_cycles\"] for e in per_program)\n", + " total_opt = sum(e[\"optimized_cycles\"] for e in per_program)\n", + " overall = {\n", + " \"n_programs\": len(per_program),\n", + " \"total_original\": total_orig,\n", + " \"total_optimized\": total_opt,\n", + " \"aggregate_reduction_pct\": (total_orig - total_opt) / total_orig if total_orig > 0 else 0.0,\n", + " \"avg_reduction_pct\": sum(e[\"reduction_pct\"] for e in per_program) / len(per_program) if per_program else 0.0,\n", + " }\n", + "\n", + " return {\n", + " \"per_program\": per_program,\n", + " \"by_level\": by_level,\n", + " \"overall\": overall,\n", + " }" + ], + "metadata": { + "id": "6LuK5ln0Cd4K" + }, + "execution_count": 47, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for held-out test set ===\n", + "\n", + "# Test 1: deterministic β€” same call returns same programs\n", + "ts1 = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "ts2 = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "\n", + "# Compare instruction lists across the two calls\n", + "all_match = all(\n", + " a[\"program\"][\"instructions\"] == b[\"program\"][\"instructions\"]\n", + " for a, b in zip(ts1, ts2)\n", + ")\n", + "print(f\"Test 1 (deterministic test set): {'PASS' if all_match else 'FAIL'}\")\n", + "\n", + "# Test 2: seeds are in held-out range\n", + "ts = generate_test_set(n_per_level=5, levels=[1, 2])\n", + "all_holdout = all(HOLDOUT_SEED_MIN <= e[\"seed\"] <= HOLDOUT_SEED_MAX for e in ts)\n", + "print(f\"Test 2 (all seeds in held-out range): {'PASS' if all_holdout else 'FAIL'}\")\n", + "\n", + "# Test 3: distribution across levels\n", + "print(f\"\\nTest 3 (test set composition for n_per_level=5, levels=[1,2]):\")\n", + "level_counts = {}\n", + "for e in ts: level_counts[e[\"level\"]] = level_counts.get(e[\"level\"], 0) + 1\n", + "print(f\" level counts: {level_counts}\")\n", + "print(f\" total: {len(ts)}\")\n", + "\n", + "# Test 4: full baseline run on a real test set (20 per level)\n", + "print(f\"\\n\" + \"=\"*60)\n", + "print(f\"Test 4 (baseline on full held-out test set, 20/level x 2 levels = 40 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "print(f\"\\n === BY LEVEL ===\")\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "print(f\"\\n === OVERALL ===\")\n", + "o = results[\"overall\"]\n", + "print(f\" n_programs : {o['n_programs']}\")\n", + "print(f\" total original cycles : {o['total_original']}\")\n", + "print(f\" total optimized cycles : {o['total_optimized']}\")\n", + "print(f\" aggregate reduction pct : {o['aggregate_reduction_pct']:.1%}\")\n", + "print(f\" avg per-program redux : {o['avg_reduction_pct']:.1%}\")\n", + "\n", + "# Test 5: per-program detail (first 3 entries) β€” useful for debugging\n", + "print(f\"\\n === PER-PROGRAM (first 3) ===\")\n", + "for entry in results[\"per_program\"][:3]:\n", + " print(f\" Level {entry['level']}, seed {entry['seed']}: \"\n", + " f\"{entry['original_cycles']} -> {entry['optimized_cycles']} \"\n", + " f\"({entry['reduction_pct']:.1%}, {entry['iterations']} iter)\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "bHmrTiFsFMsR", + "outputId": "2a7864a9-d601-491b-ec92-d6fd075a02bd" + }, + "execution_count": 48, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (deterministic test set): PASS\n", + "Test 2 (all seeds in held-out range): PASS\n", + "\n", + "Test 3 (test set composition for n_per_level=5, levels=[1,2]):\n", + " level counts: {1: 5, 2: 5}\n", + " total: 10\n", + "\n", + "============================================================\n", + "Test 4 (baseline on full held-out test set, 20/level x 2 levels = 40 programs):\n", + "\n", + " === BY LEVEL ===\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + "\n", + " === OVERALL ===\n", + " n_programs : 40\n", + " total original cycles : 703\n", + " total optimized cycles : 329\n", + " aggregate reduction pct : 53.2%\n", + " avg per-program redux : 50.8%\n", + "\n", + " === PER-PROGRAM (first 3) ===\n", + " Level 1, seed 10000: 11 -> 5 (54.5%, 2 iter)\n", + " Level 1, seed 10001: 10 -> 5 (50.0%, 2 iter)\n", + " Level 1, seed 10002: 10 -> 5 (50.0%, 2 iter)\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_3():\n", + " \"\"\"\n", + " Generate a Level 3 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 15-20 instructions\n", + " - 5-7 CONSTs\n", + " - 2 LOADs (both forced to flow into final STORE)\n", + " - 6-9 arithmetic ops with deeper chaining\n", + " - 2-4 dead variables sprinkled throughout\n", + " - 1 main STORE (optionally a secondary STORE)\n", + " - DIV always uses literal non-zero divisor (Design Assumption #4)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 5-7 CONSTs\n", + " num_consts = random.randint(5, 7)\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 15)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + "\n", + " # Step 2: 2 LOADs from initial memory\n", + " loaded_vars = []\n", + " for i in range(2):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: 6-9 arithmetic ops, chained\n", + " num_arith = random.randint(6, 9)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " src1 = random.choice(available_vars)\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3])\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 4: 2-4 dead CONSTs scattered throughout\n", + " num_dead = random.randint(2, 4)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 25)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 5: Force BOTH LOADed variables into the final result chain.\n", + " # Combine last_result with both loaded values via two ADDs.\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": after_load1,\n", + " \"src1\": last_result,\n", + " \"src2\": loaded_vars[0],\n", + " })\n", + "\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\",\n", + " \"dest\": after_load2,\n", + " \"src1\": after_load1,\n", + " \"src2\": loaded_vars[1],\n", + " })\n", + " last_result = after_load2\n", + "\n", + " # Step 6: Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " # Step 7 (optional, 40% chance): Secondary STORE to a different address.\n", + " # Stores an intermediate constant β€” gives the agent another observable to preserve.\n", + " has_secondary_store = random.random() < 0.4\n", + " observable_addrs = [0]\n", + " if has_secondary_store:\n", + " # Pick one of the early-defined CONSTs as the secondary value\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr1\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(1)\n", + "\n", + " # Step 8: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr_in0\": 2,\n", + " \"addr_in1\": 3,\n", + " }\n", + " initial_mem = {\n", + " 2: random.randint(1, 25),\n", + " 3: random.randint(1, 25),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 3 with the test set generator\n", + "LEVEL_GENERATORS[3] = generate_level_3" + ], + "metadata": { + "id": "LtziAqGJFPN4" + }, + "execution_count": 49, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 3 generator ===\n", + "\n", + "# Test 1: spot-check three seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_3()\n", + " print(f\"\\n=== Level 3, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + " print(f\"observable_addrs : {prog['observable_addrs']}\")\n", + " print(f\"initial_vars : {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + "\n", + "# Test 2: Full pipeline on Level 3, seed=42\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (-O3 baseline on Level 3, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_3()\n", + "print(f\"Original ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "result = evaluate_baseline(prog)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" cycle_reduction : {result['cycle_reduction']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "print(f\"\\nOptimized:\")\n", + "print(dump_ir({**prog, \"instructions\": result[\"optimized_instructions\"]}))\n", + "\n", + "# Test 3: aggregate stats across 20 Level 3 programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 3 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(200, 220):\n", + " random.seed(seed)\n", + " prog = generate_level_3()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 4: held-out test set with Level 3 included\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (held-out test set, Levels 1-3, 20 each):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n Overall: n={o['n_programs']}, \"\n", + " f\"agg reduction={o['aggregate_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "7Edy1LM1TJzg", + "outputId": "52e99889-28ac-4bd4-ea42-f0f09ee23ac5" + }, + "execution_count": 50, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 3, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v4 = 4 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 3 # CONST [1 cycle]\n", + "7: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v9 = v0 + v0 # ADD [1 cycle]\n", + "10: v10 = v8 + 0 # ADD [1 cycle]\n", + "11: v16 = 4 # CONST [1 cycle]\n", + "12: v11 = v8 - v6 # SUB [1 cycle]\n", + "13: v15 = 3 # CONST [1 cycle]\n", + "14: v12 = v4 - v0 # SUB [1 cycle]\n", + "15: v13 = v5 - v4 # SUB [1 cycle]\n", + "16: v14 = v12 - 2 # SUB [1 cycle]\n", + "17: v17 = v14 + v7 # ADD [1 cycle]\n", + "18: v18 = v17 + v8 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 29 cycles\n", + "observable_addrs : [0]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 20, 3: 9}\n", + "\n", + "=== Level 3, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v15 = 10 # CONST [1 cycle]\n", + "1: v0 = 10 # CONST [1 cycle]\n", + "2: v1 = 14 # CONST [1 cycle]\n", + "3: v16 = 14 # CONST [1 cycle]\n", + "4: v2 = 13 # CONST [1 cycle]\n", + "5: v3 = 13 # CONST [1 cycle]\n", + "6: v4 = 2 # CONST [1 cycle]\n", + "7: v5 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v6 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v17 = 6 # CONST [1 cycle]\n", + "10: v7 = v3 + v3 # ADD [1 cycle]\n", + "11: v8 = v1 // 1 # DIV [5 cycles]\n", + "12: v9 = v0 // 3 # DIV [5 cycles]\n", + "13: v10 = v1 - v5 # SUB [1 cycle]\n", + "14: v11 = v10 + 0 # ADD [1 cycle]\n", + "15: v12 = v6 // 1 # DIV [5 cycles]\n", + "16: v13 = v7 - v8 # SUB [1 cycle]\n", + "17: v14 = v10 - v3 # SUB [1 cycle]\n", + "18: v18 = v14 + v5 # ADD [1 cycle]\n", + "19: v19 = v18 + v6 # ADD [1 cycle]\n", + "20: mem[addr0] = v19 # STORE [4 cycles]\n", + "21: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 46 cycles\n", + "observable_addrs : [0, 1]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 24, 3: 23}\n", + "\n", + "=== Level 3, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 3 # CONST [1 cycle]\n", + "1: v1 = 7 # CONST [1 cycle]\n", + "2: v2 = 11 # CONST [1 cycle]\n", + "3: v3 = 1 # CONST [1 cycle]\n", + "4: v4 = 2 # CONST [1 cycle]\n", + "5: v16 = 18 # CONST [1 cycle]\n", + "6: v5 = 14 # CONST [1 cycle]\n", + "7: v6 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v7 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v14 = 5 # CONST [1 cycle]\n", + "10: v15 = 4 # CONST [1 cycle]\n", + "11: v8 = v3 * v0 # MUL [3 cycles]\n", + "12: v9 = v1 + v3 # ADD [1 cycle]\n", + "13: v10 = v0 + v9 # ADD [1 cycle]\n", + "14: v11 = v10 + v10 # ADD [1 cycle]\n", + "15: v12 = v6 + v0 # ADD [1 cycle]\n", + "16: v13 = v2 - 2 # SUB [1 cycle]\n", + "17: v17 = v13 + v6 # ADD [1 cycle]\n", + "18: v18 = v17 + v7 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "20: mem[addr1] = v4 # STORE [4 cycles]\n", + "// TOTAL: 35 cycles\n", + "observable_addrs : [0, 1]\n", + "initial_vars : {'addr0': 0, 'addr1': 1, 'addr_in0': 2, 'addr_in1': 3}\n", + "initial_mem : {2: 21, 3: 7}\n", + "\n", + "============================================================\n", + "Test 2 (-O3 baseline on Level 3, seed=42):\n", + "Original (29 cycles):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v4 = 4 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 3 # CONST [1 cycle]\n", + "7: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v9 = v0 + v0 # ADD [1 cycle]\n", + "10: v10 = v8 + 0 # ADD [1 cycle]\n", + "11: v16 = 4 # CONST [1 cycle]\n", + "12: v11 = v8 - v6 # SUB [1 cycle]\n", + "13: v15 = 3 # CONST [1 cycle]\n", + "14: v12 = v4 - v0 # SUB [1 cycle]\n", + "15: v13 = v5 - v4 # SUB [1 cycle]\n", + "16: v14 = v12 - 2 # SUB [1 cycle]\n", + "17: v17 = v14 + v7 # ADD [1 cycle]\n", + "18: v18 = v17 + v8 # ADD [1 cycle]\n", + "19: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 29 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 29\n", + " optimized_cycles : 14\n", + " cycle_reduction : 15\n", + " reduction_pct : 51.7%\n", + " iterations : 2\n", + "\n", + "Optimized:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "1: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "2: v17 = 0 + v7 # ADD [1 cycle]\n", + "3: v18 = v17 + v8 # ADD [1 cycle]\n", + "4: mem[addr0] = v18 # STORE [4 cycles]\n", + "// TOTAL: 14 cycles\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 3 programs):\n", + " total original cycles : 839\n", + " total optimized cycles : 329\n", + " avg reduction pct : 60.8%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 4 (held-out test set, Levels 1-3, 20 each):\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig=43.4, avg opt=17.6, agg reduction=59.5%\n", + "\n", + " Overall: n=60, agg reduction=56.7%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_4():\n", + " \"\"\"\n", + " Generate a Level 4 Toy-IR program with fan-out topology.\n", + "\n", + " Distinct from Levels 1-3: 2-3 'hub' variables fan out to drive\n", + " most arithmetic operations. A single optimization at a hub cascades\n", + " through many dependents, giving the agent measurable leverage from\n", + " correctly identifying high-value optimization targets.\n", + "\n", + " Characteristics:\n", + " - 20-28 instructions\n", + " - 5-7 CONSTs (some become hubs)\n", + " - 2 LOADs (one usually a hub)\n", + " - 10-14 arithmetic ops, ~70% using a hub variable as src1\n", + " - 3-5 dead variables\n", + " - 1-2 STOREs (both LOADs flow into final result)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 5-7 CONSTs\n", + " num_consts = random.randint(5, 7)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 15)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: 2 LOADs\n", + " loaded_vars = []\n", + " for i in range(2):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: Designate 2-3 hub variables.\n", + " # Hubs come from a mix of constants and loads to ensure CF can exploit\n", + " # constant hubs while non-constant (LOAD) hubs force the agent to reason\n", + " # about partially-resolvable structure.\n", + " num_hubs = random.randint(2, 3)\n", + " candidate_hubs = const_vars[:3] + loaded_vars # bias toward early-defined vars\n", + " hub_vars = random.sample(candidate_hubs, min(num_hubs, len(candidate_hubs)))\n", + "\n", + " # Step 4: 10-14 arithmetic ops, ~70% using a hub as src1\n", + " num_arith = random.randint(10, 14)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " # 70% chance src1 is a hub (drives the fan-out structure)\n", + " if random.random() < 0.70:\n", + " src1 = random.choice(hub_vars)\n", + " else:\n", + " src1 = random.choice(available_vars)\n", + "\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3])\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 5: 3-5 dead CONSTs scattered throughout\n", + " num_dead = random.randint(3, 5)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 30)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 6: Force both LOADed values into the final chain\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load1,\n", + " \"src1\": last_result, \"src2\": loaded_vars[0],\n", + " })\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load2,\n", + " \"src1\": after_load1, \"src2\": loaded_vars[1],\n", + " })\n", + " last_result = after_load2\n", + "\n", + " # Step 7: Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " # Step 8 (40% chance): secondary STORE\n", + " has_secondary_store = random.random() < 0.4\n", + " observable_addrs = [0]\n", + " if has_secondary_store:\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr1\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(1)\n", + "\n", + " # Step 9: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr_in0\": 2,\n", + " \"addr_in1\": 3,\n", + " }\n", + " initial_mem = {\n", + " 2: random.randint(1, 30),\n", + " 3: random.randint(1, 30),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 4\n", + "LEVEL_GENERATORS[4] = generate_level_4" + ], + "metadata": { + "id": "cyi06K_mVhUf" + }, + "execution_count": 51, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 4 generator ===\n", + "\n", + "# Test 1: spot-check three seeds β€” verify fan-out is visible\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_4()\n", + " print(f\"\\n=== Level 4, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + "\n", + "# Test 2: structural verification β€” count fan-out\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (verify fan-out structure):\")\n", + "random.seed(42)\n", + "prog = generate_level_4()\n", + "\n", + "# Count how often each variable is used as src1 or src2\n", + "usage_count = {}\n", + "for instr in prog[\"instructions\"]:\n", + " for src_field in (\"src1\", \"src2\"):\n", + " src = instr[src_field]\n", + " if isinstance(src, str):\n", + " usage_count[src] = usage_count.get(src, 0) + 1\n", + "\n", + "# Sort by usage\n", + "sorted_usage = sorted(usage_count.items(), key=lambda x: -x[1])\n", + "print(f\" Top 5 most-used variables (the 'hubs'):\")\n", + "for var, count in sorted_usage[:5]:\n", + " print(f\" {var}: used {count} times\")\n", + "print(f\" (In linear Level 3, top variables typically used 1-3 times.)\")\n", + "\n", + "# Test 3: baseline on Level 4\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 baseline on Level 4, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_4()\n", + "print(f\"Original ({count_cycles(prog['instructions'])} cycles):\")\n", + "print(dump_ir(prog))\n", + "\n", + "result = evaluate_baseline(prog)\n", + "print(f\"\\nBaseline result:\")\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "\n", + "# Test 4: aggregate stats on Level 4\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (-O3 across 20 random Level 4 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(300, 320):\n", + " random.seed(seed)\n", + " prog = generate_level_4()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 5: full held-out test set (Levels 1-4)\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 5 (held-out test set, Levels 1-4, 20 each = 80 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3, 4])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n Overall: n={o['n_programs']}, agg reduction={o['aggregate_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "m8yNd1VYVhuL", + "outputId": "0ed42054-3ebb-4e5e-bae7-267a1ad9dd9c" + }, + "execution_count": 52, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 4, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v25 = 15 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = 4 # CONST [1 cycle]\n", + "7: v6 = 3 # CONST [1 cycle]\n", + "8: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "10: v9 = v8 // 5 # DIV [5 cycles]\n", + "11: v10 = v8 + v6 # ADD [1 cycle]\n", + "12: v11 = v8 - v2 # SUB [1 cycle]\n", + "13: v12 = v0 // 1 # DIV [5 cycles]\n", + "14: v13 = v9 // 3 # DIV [5 cycles]\n", + "15: v14 = v0 + v1 # ADD [1 cycle]\n", + "16: v23 = 21 # CONST [1 cycle]\n", + "17: v15 = v0 * v9 # MUL [3 cycles]\n", + "18: v16 = v8 - v9 # SUB [1 cycle]\n", + "19: v17 = v12 + v8 # ADD [1 cycle]\n", + "20: v18 = v0 // 2 # DIV [5 cycles]\n", + "21: v24 = 13 # CONST [1 cycle]\n", + "22: v19 = v8 * v5 # MUL [3 cycles]\n", + "23: v20 = v8 - 2 # SUB [1 cycle]\n", + "24: v21 = v0 + 3 # ADD [1 cycle]\n", + "25: v22 = v18 * 2 # MUL [3 cycles]\n", + "26: v26 = v22 + v7 # ADD [1 cycle]\n", + "27: v27 = v26 + v8 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 64 cycles\n", + "\n", + "=== Level 4, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 14 # CONST [1 cycle]\n", + "2: v2 = 13 # CONST [1 cycle]\n", + "3: v3 = 13 # CONST [1 cycle]\n", + "4: v4 = 2 # CONST [1 cycle]\n", + "5: v21 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in0] # LOAD [4 cycles]\n", + "7: v6 = mem[addr_in1] # LOAD [4 cycles]\n", + "8: v7 = v0 // 4 # DIV [5 cycles]\n", + "9: v8 = v7 // 2 # DIV [5 cycles]\n", + "10: v9 = v0 + v8 # ADD [1 cycle]\n", + "11: v10 = v0 + v6 # ADD [1 cycle]\n", + "12: v22 = 17 # CONST [1 cycle]\n", + "13: v11 = v7 + v8 # ADD [1 cycle]\n", + "14: v12 = v5 - v4 # SUB [1 cycle]\n", + "15: v23 = 12 # CONST [1 cycle]\n", + "16: v13 = v10 + v1 # ADD [1 cycle]\n", + "17: v14 = v4 - v1 # SUB [1 cycle]\n", + "18: v15 = v8 * v14 # MUL [3 cycles]\n", + "19: v20 = 22 # CONST [1 cycle]\n", + "20: v16 = v6 // 3 # DIV [5 cycles]\n", + "21: v17 = v1 // 4 # DIV [5 cycles]\n", + "22: v18 = v1 - v5 # SUB [1 cycle]\n", + "23: v19 = v11 * v2 # MUL [3 cycles]\n", + "24: v24 = v19 + v5 # ADD [1 cycle]\n", + "25: v25 = v24 + v6 # ADD [1 cycle]\n", + "26: mem[addr0] = v25 # STORE [4 cycles]\n", + "// TOTAL: 56 cycles\n", + "\n", + "=== Level 4, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v24 = 12 # CONST [1 cycle]\n", + "1: v0 = 3 # CONST [1 cycle]\n", + "2: v1 = 7 # CONST [1 cycle]\n", + "3: v2 = 11 # CONST [1 cycle]\n", + "4: v3 = 1 # CONST [1 cycle]\n", + "5: v4 = 2 # CONST [1 cycle]\n", + "6: v5 = 14 # CONST [1 cycle]\n", + "7: v6 = mem[addr_in0] # LOAD [4 cycles]\n", + "8: v7 = mem[addr_in1] # LOAD [4 cycles]\n", + "9: v8 = v2 - 1 # SUB [1 cycle]\n", + "10: v9 = v2 + v3 # ADD [1 cycle]\n", + "11: v25 = 15 # CONST [1 cycle]\n", + "12: v10 = v2 + v0 # ADD [1 cycle]\n", + "13: v11 = v2 - 2 # SUB [1 cycle]\n", + "14: v22 = 23 # CONST [1 cycle]\n", + "15: v12 = v2 - 2 # SUB [1 cycle]\n", + "16: v13 = v2 + v9 # ADD [1 cycle]\n", + "17: v14 = v0 - v7 # SUB [1 cycle]\n", + "18: v15 = v2 // 2 # DIV [5 cycles]\n", + "19: v16 = v0 + v14 # ADD [1 cycle]\n", + "20: v17 = v0 * v5 # MUL [3 cycles]\n", + "21: v18 = v2 * 0 # MUL [3 cycles]\n", + "22: v19 = v0 * v18 # MUL [3 cycles]\n", + "23: v23 = 29 # CONST [1 cycle]\n", + "24: v20 = v0 // 4 # DIV [5 cycles]\n", + "25: v21 = v20 + 3 # ADD [1 cycle]\n", + "26: v26 = v21 + v6 # ADD [1 cycle]\n", + "27: v27 = v26 + v7 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v24 # STORE [4 cycles]\n", + "// TOTAL: 56 cycles\n", + "\n", + "============================================================\n", + "Test 2 (verify fan-out structure):\n", + " Top 5 most-used variables (the 'hubs'):\n", + " v8: used 8 times\n", + " v0: used 5 times\n", + " v9: used 3 times\n", + " v1: used 2 times\n", + " addr_in0: used 1 times\n", + " (In linear Level 3, top variables typically used 1-3 times.)\n", + "\n", + "============================================================\n", + "Test 3 (-O3 baseline on Level 4, seed=42):\n", + "Original (64 cycles):\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 12 # CONST [1 cycle]\n", + "3: v3 = 5 # CONST [1 cycle]\n", + "4: v25 = 15 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = 4 # CONST [1 cycle]\n", + "7: v6 = 3 # CONST [1 cycle]\n", + "8: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "10: v9 = v8 // 5 # DIV [5 cycles]\n", + "11: v10 = v8 + v6 # ADD [1 cycle]\n", + "12: v11 = v8 - v2 # SUB [1 cycle]\n", + "13: v12 = v0 // 1 # DIV [5 cycles]\n", + "14: v13 = v9 // 3 # DIV [5 cycles]\n", + "15: v14 = v0 + v1 # ADD [1 cycle]\n", + "16: v23 = 21 # CONST [1 cycle]\n", + "17: v15 = v0 * v9 # MUL [3 cycles]\n", + "18: v16 = v8 - v9 # SUB [1 cycle]\n", + "19: v17 = v12 + v8 # ADD [1 cycle]\n", + "20: v18 = v0 // 2 # DIV [5 cycles]\n", + "21: v24 = 13 # CONST [1 cycle]\n", + "22: v19 = v8 * v5 # MUL [3 cycles]\n", + "23: v20 = v8 - 2 # SUB [1 cycle]\n", + "24: v21 = v0 + 3 # ADD [1 cycle]\n", + "25: v22 = v18 * 2 # MUL [3 cycles]\n", + "26: v26 = v22 + v7 # ADD [1 cycle]\n", + "27: v27 = v26 + v8 # ADD [1 cycle]\n", + "28: mem[addr0] = v27 # STORE [4 cycles]\n", + "29: mem[addr1] = v1 # STORE [4 cycles]\n", + "// TOTAL: 64 cycles\n", + "\n", + "Baseline result:\n", + " original_cycles : 64\n", + " optimized_cycles : 19\n", + " reduction_pct : 70.3%\n", + " iterations : 2\n", + "\n", + "============================================================\n", + "Test 4 (-O3 across 20 random Level 4 programs):\n", + " total original cycles : 1102\n", + " total optimized cycles : 345\n", + " avg reduction pct : 68.7%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 5 (held-out test set, Levels 1-4, 20 each = 80 programs):\n", + " Level 1: n=20, avg orig=9.9, avg opt=5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig=25.2, avg opt=11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig=43.4, avg opt=17.6, agg reduction=59.5%\n", + " Level 4: n=20, avg orig=55.0, avg opt=17.6, agg reduction=68.0%\n", + "\n", + " Overall: n=80, agg reduction=61.4%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_5():\n", + " \"\"\"\n", + " Generate a Level 5 Toy-IR program β€” maximum adversarial density.\n", + "\n", + " Level 5 is the hardest curriculum tier. Programs are long (30-40\n", + " instructions), peppered with dead code, heavy on expensive ops\n", + " (MUL/DIV), and feature multiple LOADs and STOREs. The 70% literal-bait\n", + " rate ensures peephole has many opportunities, and the deep arithmetic\n", + " chains create real differentiation between optimization strategies.\n", + "\n", + " Distinct from Level 4: longer, denser, more peephole bait, more\n", + " expensive operations, multiple observable outputs. The absolute cycle\n", + " savings are huge, magnifying any sub-optimal pass orderings.\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = []\n", + "\n", + " # Step 1: 7-9 CONSTs\n", + " num_consts = random.randint(7, 9)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 20)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: 3 LOADs\n", + " loaded_vars = []\n", + " for i in range(3):\n", + " addr_name = f\"addr_in{i}\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": addr_name, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + " loaded_vars.append(loaded_var)\n", + "\n", + " # Step 3: Designate 3-4 hub variables\n", + " num_hubs = random.randint(3, 4)\n", + " candidate_hubs = const_vars[:5] + loaded_vars\n", + " hub_vars = random.sample(candidate_hubs, min(num_hubs, len(candidate_hubs)))\n", + "\n", + " # Step 4: 15-22 arithmetic ops, heavy on expensive ops + peephole bait\n", + " num_arith = random.randint(15, 22)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " # Bias toward expensive ops (MUL/DIV) for higher cycle savings on optimization\n", + " op = random.choices(\n", + " [\"ADD\", \"SUB\", \"MUL\", \"DIV\"],\n", + " weights=[2, 2, 3, 2], # MUL slightly preferred\n", + " )[0]\n", + "\n", + " # 70% chance of literal src2 β†’ peephole bait\n", + " use_literal_src2 = random.random() < 0.70\n", + "\n", + " # 60% chance src1 is a hub\n", + " if random.random() < 0.60:\n", + " src1 = random.choice(hub_vars)\n", + " else:\n", + " src1 = random.choice(available_vars)\n", + "\n", + " if use_literal_src2:\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5]) # peephole bait: 1\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3, 4]) # peephole bait: 0, 1, 2\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4, 5])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 5: 5-8 dead CONSTs\n", + " num_dead = random.randint(5, 8)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 50)\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + "\n", + " # Step 6: Force all 3 LOADed values into the final chain\n", + " after_load1 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load1,\n", + " \"src1\": last_result, \"src2\": loaded_vars[0],\n", + " })\n", + " after_load2 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load2,\n", + " \"src1\": after_load1, \"src2\": loaded_vars[1],\n", + " })\n", + " after_load3 = new_var()\n", + " instructions.append({\n", + " \"op\": \"ADD\", \"dest\": after_load3,\n", + " \"src1\": after_load2, \"src2\": loaded_vars[2],\n", + " })\n", + " last_result = after_load3\n", + "\n", + " # Step 7: Multiple STOREs (always 2-3)\n", + " num_stores = random.randint(2, 3)\n", + "\n", + " # Main STORE\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + " observable_addrs = [0]\n", + "\n", + " # Secondary STOREs use early CONSTs (gives DCE more variety to handle)\n", + " early_const_vars = [\n", + " instr[\"dest\"] for instr in instructions[:num_consts]\n", + " if instr[\"op\"] == \"CONST\"\n", + " ]\n", + " for i in range(1, num_stores):\n", + " if early_const_vars:\n", + " secondary_src = random.choice(early_const_vars)\n", + " addr_idx = i # addr1, addr2\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": f\"addr{addr_idx}\", \"src1\": secondary_src, \"src2\": None,\n", + " })\n", + " observable_addrs.append(addr_idx)\n", + "\n", + " # Step 8: Initial state\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr1\": 1,\n", + " \"addr2\": 2,\n", + " \"addr_in0\": 3,\n", + " \"addr_in1\": 4,\n", + " \"addr_in2\": 5,\n", + " }\n", + " initial_mem = {\n", + " 3: random.randint(1, 50),\n", + " 4: random.randint(1, 50),\n", + " 5: random.randint(1, 50),\n", + " }\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": observable_addrs,\n", + " }\n", + "\n", + "\n", + "# Register Level 5\n", + "LEVEL_GENERATORS[5] = generate_level_5" + ], + "metadata": { + "id": "vExn6K85Vk7D" + }, + "execution_count": 53, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 5 generator ===\n", + "\n", + "# Test 1: spot-check three seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_5()\n", + " print(f\"\\n=== Level 5, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + "\n", + "# Test 2: baseline on Level 5\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 2 (-O3 baseline on Level 5, seed=42):\")\n", + "random.seed(42)\n", + "prog = generate_level_5()\n", + "result = evaluate_baseline(prog)\n", + "print(f\" original_cycles : {result['original_cycles']}\")\n", + "print(f\" optimized_cycles : {result['optimized_cycles']}\")\n", + "print(f\" reduction_pct : {result['reduction_pct']:.1%}\")\n", + "print(f\" iterations : {result['iterations']}\")\n", + "\n", + "# Test 3: aggregate on Level 5\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 3 (-O3 across 20 random Level 5 programs):\")\n", + "total_orig = 0\n", + "total_opt = 0\n", + "iter_dist = {}\n", + "for seed in range(400, 420):\n", + " random.seed(seed)\n", + " prog = generate_level_5()\n", + " result = evaluate_baseline(prog)\n", + " total_orig += result[\"original_cycles\"]\n", + " total_opt += result[\"optimized_cycles\"]\n", + " iter_dist[result[\"iterations\"]] = iter_dist.get(result[\"iterations\"], 0) + 1\n", + "\n", + "avg_pct = (total_orig - total_opt) / total_orig\n", + "print(f\" total original cycles : {total_orig}\")\n", + "print(f\" total optimized cycles : {total_opt}\")\n", + "print(f\" avg reduction pct : {avg_pct:.1%}\")\n", + "print(f\" iteration distribution : {iter_dist}\")\n", + "\n", + "# Test 4: full headline plot β€” held-out test set, all 5 levels, 20 each = 100 programs\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\"Test 4 (FULL HELD-OUT TEST SET, Levels 1-5, 100 programs):\")\n", + "test_set = generate_test_set(n_per_level=20, levels=[1, 2, 3, 4, 5])\n", + "results = run_baseline_on_test_set(test_set)\n", + "\n", + "print(f\"\\n === BY LEVEL ===\")\n", + "for lvl in sorted(results[\"by_level\"].keys()):\n", + " s = results[\"by_level\"][lvl]\n", + " print(f\" Level {lvl}: n={s['n_programs']}, \"\n", + " f\"avg orig={s['avg_original']:6.1f}, \"\n", + " f\"avg opt={s['avg_optimized']:6.1f}, \"\n", + " f\"agg reduction={s['aggregate_reduction_pct']:.1%}\")\n", + "\n", + "o = results[\"overall\"]\n", + "print(f\"\\n === OVERALL ===\")\n", + "print(f\" n_programs : {o['n_programs']}\")\n", + "print(f\" total original cycles : {o['total_original']}\")\n", + "print(f\" total optimized cycles : {o['total_optimized']}\")\n", + "print(f\" aggregate reduction pct : {o['aggregate_reduction_pct']:.1%}\")\n", + "print(f\" avg per-program redux : {o['avg_reduction_pct']:.1%}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Igt72oqfWuPV", + "outputId": "24d886eb-76a0-4e4c-a665-ed8e804fa8db" + }, + "execution_count": 54, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 5, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1], mem[2]\n", + "0: v32 = 36 # CONST [1 cycle]\n", + "1: v0 = 4 # CONST [1 cycle]\n", + "2: v1 = 1 # CONST [1 cycle]\n", + "3: v2 = 9 # CONST [1 cycle]\n", + "4: v3 = 8 # CONST [1 cycle]\n", + "5: v4 = 8 # CONST [1 cycle]\n", + "6: v5 = 5 # CONST [1 cycle]\n", + "7: v33 = 44 # CONST [1 cycle]\n", + "8: v6 = 4 # CONST [1 cycle]\n", + "9: v7 = 18 # CONST [1 cycle]\n", + "10: v8 = 3 # CONST [1 cycle]\n", + "11: v9 = mem[addr_in0] # LOAD [4 cycles]\n", + "12: v10 = mem[addr_in1] # LOAD [4 cycles]\n", + "13: v11 = mem[addr_in2] # LOAD [4 cycles]\n", + "14: v12 = v1 * 1 # MUL [3 cycles]\n", + "15: v13 = v12 * 1 # MUL [3 cycles]\n", + "16: v14 = v10 * 0 # MUL [3 cycles]\n", + "17: v35 = 49 # CONST [1 cycle]\n", + "18: v37 = 8 # CONST [1 cycle]\n", + "19: v15 = v9 + 2 # ADD [1 cycle]\n", + "20: v31 = 30 # CONST [1 cycle]\n", + "21: v16 = v1 // 5 # DIV [5 cycles]\n", + "22: v36 = 50 # CONST [1 cycle]\n", + "23: v17 = v11 - 4 # SUB [1 cycle]\n", + "24: v18 = v9 + 0 # ADD [1 cycle]\n", + "25: v30 = 25 # CONST [1 cycle]\n", + "26: v19 = v1 // 2 # DIV [5 cycles]\n", + "27: v20 = v0 - 4 # SUB [1 cycle]\n", + "28: v21 = v1 * 3 # MUL [3 cycles]\n", + "29: v22 = v7 - v21 # SUB [1 cycle]\n", + "30: v23 = v0 - v10 # SUB [1 cycle]\n", + "31: v24 = v18 - 2 # SUB [1 cycle]\n", + "32: v25 = v20 + 3 # ADD [1 cycle]\n", + "33: v26 = v17 + 2 # ADD [1 cycle]\n", + "34: v27 = v10 * 1 # MUL [3 cycles]\n", + "35: v28 = v0 // 1 # DIV [5 cycles]\n", + "36: v29 = v19 + 0 # ADD [1 cycle]\n", + "37: v34 = 44 # CONST [1 cycle]\n", + "38: v38 = v29 + v9 # ADD [1 cycle]\n", + "39: v39 = v38 + v10 # ADD [1 cycle]\n", + "40: v40 = v39 + v11 # ADD [1 cycle]\n", + "41: mem[addr0] = v40 # STORE [4 cycles]\n", + "42: mem[addr1] = v1 # STORE [4 cycles]\n", + "43: mem[addr2] = v33 # STORE [4 cycles]\n", + "// TOTAL: 84 cycles\n", + "\n", + "=== Level 5, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: v34 = 29 # CONST [1 cycle]\n", + "1: v0 = 19 # CONST [1 cycle]\n", + "2: v33 = 6 # CONST [1 cycle]\n", + "3: v1 = 3 # CONST [1 cycle]\n", + "4: v2 = 9 # CONST [1 cycle]\n", + "5: v3 = 4 # CONST [1 cycle]\n", + "6: v32 = 3 # CONST [1 cycle]\n", + "7: v4 = 16 # CONST [1 cycle]\n", + "8: v5 = 15 # CONST [1 cycle]\n", + "9: v6 = 16 # CONST [1 cycle]\n", + "10: v7 = mem[addr_in0] # LOAD [4 cycles]\n", + "11: v8 = mem[addr_in1] # LOAD [4 cycles]\n", + "12: v9 = mem[addr_in2] # LOAD [4 cycles]\n", + "13: v10 = v8 - v4 # SUB [1 cycle]\n", + "14: v11 = v5 * 0 # MUL [3 cycles]\n", + "15: v12 = v8 + 1 # ADD [1 cycle]\n", + "16: v13 = v8 // 5 # DIV [5 cycles]\n", + "17: v36 = 16 # CONST [1 cycle]\n", + "18: v35 = 49 # CONST [1 cycle]\n", + "19: v14 = v8 - 2 # SUB [1 cycle]\n", + "20: v15 = v14 // 1 # DIV [5 cycles]\n", + "21: v31 = 6 # CONST [1 cycle]\n", + "22: v16 = v3 + v10 # ADD [1 cycle]\n", + "23: v17 = v8 // 2 # DIV [5 cycles]\n", + "24: v18 = v16 - 3 # SUB [1 cycle]\n", + "25: v19 = v8 * 3 # MUL [3 cycles]\n", + "26: v20 = v11 * 0 # MUL [3 cycles]\n", + "27: v21 = v16 - 3 # SUB [1 cycle]\n", + "28: v22 = v9 - v19 # SUB [1 cycle]\n", + "29: v23 = v0 * 0 # MUL [3 cycles]\n", + "30: v24 = v7 * 3 # MUL [3 cycles]\n", + "31: v25 = v8 * v8 # MUL [3 cycles]\n", + "32: v26 = v12 * 4 # MUL [3 cycles]\n", + "33: v27 = v8 // 1 # DIV [5 cycles]\n", + "34: v28 = v8 * 3 # MUL [3 cycles]\n", + "35: v29 = v9 // 4 # DIV [5 cycles]\n", + "36: v30 = v17 * v18 # MUL [3 cycles]\n", + "37: v37 = v30 + v7 # ADD [1 cycle]\n", + "38: v38 = v37 + v8 # ADD [1 cycle]\n", + "39: v39 = v38 + v9 # ADD [1 cycle]\n", + "40: mem[addr0] = v39 # STORE [4 cycles]\n", + "41: mem[addr1] = v32 # STORE [4 cycles]\n", + "// TOTAL: 95 cycles\n", + "\n", + "=== Level 5, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0], mem[1], mem[2]\n", + "0: v0 = 5 # CONST [1 cycle]\n", + "1: v1 = 13 # CONST [1 cycle]\n", + "2: v2 = 2 # CONST [1 cycle]\n", + "3: v3 = 3 # CONST [1 cycle]\n", + "4: v4 = 18 # CONST [1 cycle]\n", + "5: v5 = 4 # CONST [1 cycle]\n", + "6: v6 = 12 # CONST [1 cycle]\n", + "7: v7 = 19 # CONST [1 cycle]\n", + "8: v8 = mem[addr_in0] # LOAD [4 cycles]\n", + "9: v34 = 1 # CONST [1 cycle]\n", + "10: v9 = mem[addr_in1] # LOAD [4 cycles]\n", + "11: v10 = mem[addr_in2] # LOAD [4 cycles]\n", + "12: v32 = 32 # CONST [1 cycle]\n", + "13: v11 = v3 - 4 # SUB [1 cycle]\n", + "14: v12 = v0 + 4 # ADD [1 cycle]\n", + "15: v13 = v9 * 1 # MUL [3 cycles]\n", + "16: v14 = v0 - 4 # SUB [1 cycle]\n", + "17: v15 = v9 // 2 # DIV [5 cycles]\n", + "18: v16 = v3 - 4 # SUB [1 cycle]\n", + "19: v33 = 17 # CONST [1 cycle]\n", + "20: v17 = v0 + 3 # ADD [1 cycle]\n", + "21: v18 = v3 * 1 # MUL [3 cycles]\n", + "22: v19 = v0 + 3 # ADD [1 cycle]\n", + "23: v20 = v0 - v5 # SUB [1 cycle]\n", + "24: v21 = v3 * 0 # MUL [3 cycles]\n", + "25: v22 = v10 * 2 # MUL [3 cycles]\n", + "26: v23 = v0 * 0 # MUL [3 cycles]\n", + "27: v24 = v9 // 1 # DIV [5 cycles]\n", + "28: v25 = v9 * 3 # MUL [3 cycles]\n", + "29: v26 = v0 - 3 # SUB [1 cycle]\n", + "30: v27 = v3 - 2 # SUB [1 cycle]\n", + "31: v28 = v0 + 0 # ADD [1 cycle]\n", + "32: v29 = v3 + 3 # ADD [1 cycle]\n", + "33: v30 = v0 // 4 # DIV [5 cycles]\n", + "34: v35 = 27 # CONST [1 cycle]\n", + "35: v31 = v3 // 2 # DIV [5 cycles]\n", + "36: v36 = 24 # CONST [1 cycle]\n", + "37: v37 = v31 + v8 # ADD [1 cycle]\n", + "38: v38 = v37 + v9 # ADD [1 cycle]\n", + "39: v39 = v38 + v10 # ADD [1 cycle]\n", + "40: mem[addr0] = v39 # STORE [4 cycles]\n", + "41: mem[addr1] = v2 # STORE [4 cycles]\n", + "42: mem[addr2] = v0 # STORE [4 cycles]\n", + "// TOTAL: 89 cycles\n", + "\n", + "============================================================\n", + "Test 2 (-O3 baseline on Level 5, seed=42):\n", + " original_cycles : 84\n", + " optimized_cycles : 29\n", + " reduction_pct : 65.5%\n", + " iterations : 2\n", + "\n", + "============================================================\n", + "Test 3 (-O3 across 20 random Level 5 programs):\n", + " total original cycles : 1659\n", + " total optimized cycles : 534\n", + " avg reduction pct : 67.8%\n", + " iteration distribution : {2: 20}\n", + "\n", + "============================================================\n", + "Test 4 (FULL HELD-OUT TEST SET, Levels 1-5, 100 programs):\n", + "\n", + " === BY LEVEL ===\n", + " Level 1: n=20, avg orig= 9.9, avg opt= 5.0, agg reduction=49.7%\n", + " Level 2: n=20, avg orig= 25.2, avg opt= 11.4, agg reduction=54.6%\n", + " Level 3: n=20, avg orig= 43.4, avg opt= 17.6, agg reduction=59.5%\n", + " Level 4: n=20, avg orig= 55.0, avg opt= 17.6, agg reduction=68.0%\n", + " Level 5: n=20, avg orig= 84.0, avg opt= 26.1, agg reduction=68.9%\n", + "\n", + " === OVERALL ===\n", + " n_programs : 100\n", + " total original cycles : 4350\n", + " total optimized cycles : 1555\n", + " aggregate reduction pct : 64.3%\n", + " avg per-program redux : 59.3%\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [], + "metadata": { + "id": "iCDDLZScWvlr" + }, "execution_count": null, "outputs": [] } diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.DS_Store b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.DS_Store differ diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb new file mode 100644 index 0000000000000000000000000000000000000000..926491c97464967c001cab5b34009ed175a8a465 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb @@ -0,0 +1,1379 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "code", + "source": [ + "## Design Assumptions (do not violate)\n", + "\n", + "# 1. **DCE never eliminates STOREs.** They define program output (final mem state).\n", + "# 2. **Addresses are distinct by construction.** Generator allocates each address variable to a unique integer; no aliasing.\n", + "# 3. **CF refuses to fold DIV by zero.** If src2 == 0 on a DIV op, leave instruction unchanged.\n", + "# 4. **Generator never emits DIV by literal zero.** When DIV is generated, src2 is always a non-zero constant or a variable known to be non-zero.\n", + "# # 5. **Integer arithmetic only.** No floats anywhere β€” avoids equivalence-check precision issues.\n", + "# 6. Generator declares `observable_addrs` per program β€” verifier compares only these mem entries.\n", + "# 7. State translator annotates observable outputs at top of dump.\n", + "# 8. Reward distinguishes broken (-1000) from valid-but-worse (small negative) β€” Harshal's formula.\n", + "# 9. Multi-input verification: 3-5 random initial states, all must match.\n", + "# 10. Integer division uses Python floor division (//). Aarush's VM must match." + ], + "metadata": { + "id": "_XI5jT2Ibvrf" + }, + "execution_count": 2, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": { + "id": "j6CQ327KWXwr" + }, + "outputs": [], + "source": [ + "# === TAC Schema v1.0 (LOCKED with Role 1 / Aarush) ===\n", + "# Reverse passes deferred to stretch goal β€” not in initial action space.\n", + "\n", + "OPS = [\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\", \"STORE\", \"NOP\"]\n", + "\n", + "CYCLE_COSTS = {\n", + " \"CONST\": 1,\n", + " \"ADD\": 1,\n", + " \"SUB\": 1,\n", + " \"MUL\": 3,\n", + " \"DIV\": 5,\n", + " \"LOAD\": 4,\n", + " \"STORE\": 4,\n", + " \"NOP\": 0,\n", + "}\n", + "\n", + "# Instruction shape: {\"op\": str, \"dest\": str|None, \"src1\": Any, \"src2\": Any}\n", + "# Operands: str = variable name, int = literal constant, None = unused\n", + "#\n", + "# Op semantics:\n", + "# CONST: dest = src1 (src1 is int literal, src2 = None)\n", + "# ADD/SUB/MUL/DIV: dest = src1 OP src2 (src1, src2 are var names or int literals)\n", + "# LOAD: dest = mem[src1] (src1 is a var holding an address)\n", + "# STORE: mem[dest] = src1 (dest is a var holding an address)\n", + "# NOP: no-op (all fields None)\n", + "#\n", + "# Program output (for equivalence check) = final memory state (mem dict).\n", + "# Programs ship as: (initial_vars: dict, initial_mem: dict, instructions: list[dict])" + ] + }, + { + "cell_type": "code", + "source": [ + "import random\n", + "\n", + "def generate_level_1():\n", + " \"\"\"\n", + " Generate a Level 1 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 4-6 instructions\n", + " - 2-3 CONST ops with literal values\n", + " - 1-2 arithmetic ops on those constants (foldable by CF)\n", + " - 0-1 dead variables (killable by DCE)\n", + " - Exactly 1 STORE at the end so the program has an observable output\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " # Step 1: Generate 2-3 constant assignments\n", + " num_consts = random.randint(2, 3)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: Generate 1-2 arithmetic ops using those constants\n", + " num_arith = random.randint(1, 2)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"MUL\"])\n", + " src1 = random.choice(const_vars)\n", + " src2 = random.choice(const_vars)\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " last_result = dest\n", + " const_vars.append(dest)\n", + "\n", + " # Step 3: Optionally add 1 dead variable (50% chance)\n", + " if random.random() < 0.5:\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None})\n", + " # Note: dead_var is intentionally never used β€” DCE should catch it.\n", + "\n", + " # Step 4: Add a STORE at the end so the program has observable output\n", + " initial_vars = {\"addr0\": 0}\n", + " instructions.append({\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None})\n", + "\n", + " initial_mem = {}\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": [0], # Aarush's verifier compares only these mem entries\n", + " }\n", + "\n", + "\n", + "# Sanity-check: generate a few programs and print them\n", + "for seed in [42, 1, 7, 99]:\n", + " random.seed(seed)\n", + " prog = generate_level_1()\n", + " print(f\"\\n=== seed={seed} ===\")\n", + " print(f\"initial_vars : {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + " print(f\"observable_addrs : {prog['observable_addrs']}\")\n", + " print(f\"instructions ({len(prog['instructions'])}):\")\n", + " for i, instr in enumerate(prog['instructions']):\n", + " print(f\" {i}: {instr}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "a86vD1OyWcB9", + "outputId": "2d6d4604-6dc1-4c29-af64-e41cf5717c0d" + }, + "execution_count": 4, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== seed=42 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (4):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "=== seed=1 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 2, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v1', 'src2': 'v1'}\n", + " 3: {'op': 'MUL', 'dest': 'v3', 'src1': 'v2', 'src2': 'v1'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=7 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (6):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=99 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 4, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 10, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v0', 'src2': 'v0'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def _resolve_operand(operand, known_constants):\n", + " \"\"\"\n", + " Given a TAC operand (string variable name or int literal),\n", + " return its concrete integer value if known, or None if unknown.\n", + " \"\"\"\n", + " if isinstance(operand, int):\n", + " return operand\n", + " if isinstance(operand, str) and operand in known_constants:\n", + " return known_constants[operand]\n", + " return None\n", + "\n", + "\n", + "def _compute(op, v1, v2):\n", + " \"\"\"Compute the result of a binary arithmetic op on two known integers.\"\"\"\n", + " if op == \"ADD\": return v1 + v2\n", + " if op == \"SUB\": return v1 - v2\n", + " if op == \"MUL\": return v1 * v2\n", + " if op == \"DIV\": return v1 // v2 # floor division (locked semantics)\n", + " raise ValueError(f\"_compute called with non-arithmetic op: {op}\")\n", + "\n", + "\n", + "def constant_folding(program):\n", + " \"\"\"\n", + " Forward pass that folds constant arithmetic into CONST ops, and\n", + " propagates known constants into instruction operands.\n", + "\n", + " Behavior:\n", + " - If both sources of an arithmetic op resolve to known integers,\n", + " replaces the instruction with a CONST holding the computed result.\n", + " - If only one source is known, still substitutes that known value\n", + " into the instruction (constant propagation), enabling downstream\n", + " passes (e.g., peephole) to recognize patterns like ADD-with-0 or MUL-by-1.\n", + " - Refuses to fold DIV by zero (Design Assumption #3).\n", + " - Always returns fresh dicts; never aliases input instructions.\n", + "\n", + " Args:\n", + " program: list of TAC instruction dicts (per locked schema).\n", + "\n", + " Returns:\n", + " new list of TAC instruction dicts. Always semantics-preserving.\n", + " Never raises, never returns None.\n", + " \"\"\"\n", + " known_constants = {}\n", + " new_program = []\n", + "\n", + " for instr in program:\n", + " # Always work on a copy β€” never alias input dicts\n", + " instr = instr.copy()\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + "\n", + " if op == \"CONST\":\n", + " known_constants[dest] = instr[\"src1\"]\n", + " new_program.append(instr)\n", + "\n", + " elif op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n", + " # === Constant propagation: substitute known constants into operands ===\n", + " if isinstance(instr[\"src1\"], str) and instr[\"src1\"] in known_constants:\n", + " instr[\"src1\"] = known_constants[instr[\"src1\"]]\n", + " if isinstance(instr[\"src2\"], str) and instr[\"src2\"] in known_constants:\n", + " instr[\"src2\"] = known_constants[instr[\"src2\"]]\n", + "\n", + " # === Try to fold ===\n", + " v1 = _resolve_operand(instr[\"src1\"], known_constants)\n", + " v2 = _resolve_operand(instr[\"src2\"], known_constants)\n", + "\n", + " if v1 is not None and v2 is not None:\n", + " # Both operands are known integers\n", + " if op == \"DIV\" and v2 == 0:\n", + " # Refuse to fold DIV by zero\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + " else:\n", + " # Fold: replace with CONST\n", + " result = _compute(op, v1, v2)\n", + " new_program.append({\n", + " \"op\": \"CONST\",\n", + " \"dest\": dest,\n", + " \"src1\": result,\n", + " \"src2\": None,\n", + " })\n", + " known_constants[dest] = result\n", + " else:\n", + " # Can't fold (at least one operand unknown).\n", + " # Instruction may still have been mutated by propagation above.\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op == \"LOAD\":\n", + " # Memory reads aren't statically resolvable\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op in (\"STORE\", \"NOP\"):\n", + " # No dest tracking needed.\n", + " # Note: we COULD propagate src1 into a STORE for downstream readability,\n", + " # but the cycle cost is unchanged and the executor handles vars fine.\n", + " # Leave STORE alone β€” keeps the code minimal.\n", + " new_program.append(instr)\n", + "\n", + " else:\n", + " # Unknown op β€” defensive pass-through\n", + " new_program.append(instr)\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "CGk8Fz4CZA3t" + }, + "execution_count": 5, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for constant_folding ===\n", + "\n", + "# Test 1: simple fold β€” ADD of two CONSTs\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result1 = constant_folding(test1)\n", + "print(\"Test 1 (simple ADD fold):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 becomes CONST c = 8\n", + "\n", + "# Test 2: chained fold β€” second op uses first op's folded result\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # β†’ 8\n", + " {\"op\": \"MUL\", \"dest\": \"d\", \"src1\": \"c\", \"src2\": \"b\"}, # β†’ 8 * 5 = 40\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"d\", \"src2\": None},\n", + "]\n", + "result2 = constant_folding(test2)\n", + "print(\"\\nTest 2 (chained fold):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: c becomes CONST 8, d becomes CONST 40\n", + "\n", + "# Test 3: DIV by zero refused\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 10, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"DIV\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # 10/0 β€” must NOT fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result3 = constant_folding(test3)\n", + "print(\"\\nTest 3 (DIV by zero refused):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still DIV, not CONST)\n", + "\n", + "# Test 4: unknown source can't fold\n", + "test4 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b is unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # can't fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result4 = constant_folding(test4)\n", + "print(\"\\nTest 4 (LOAD makes b unknown, ADD not folded):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still ADD)\n", + "\n", + "# Test 5: idempotence β€” running CF on a generated program\n", + "random.seed(42)\n", + "prog = generate_level_1()\n", + "folded = constant_folding(prog[\"instructions\"])\n", + "print(\"\\nTest 5 (CF on generated Level 1 program, seed=42):\")\n", + "print(\"Before:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "print(\"After:\")\n", + "for i, instr in enumerate(folded): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 6: constant propagation β€” only one source is known\n", + "test6 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # a known (=0), b unknown\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result6 = constant_folding(test6)\n", + "print(\"\\nTest 6 (propagation: a=0 substituted into ADD even though b unknown):\")\n", + "for i, instr in enumerate(result6): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 is still ADD (can't fold β€” b unknown), but src1 is now literal 0, not 'a'\n", + "# This sets up peephole to recognize \"ADD with 0\" later" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "WLKSTVa8fRF8", + "outputId": "e9480e1e-5d59-4afe-a884-bae0fa4a2192" + }, + "execution_count": 6, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (simple ADD fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 2 (chained fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'd', 'src1': 40, 'src2': None}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'd', 'src2': None}\n", + "\n", + "Test 3 (DIV by zero refused):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'DIV', 'dest': 'c', 'src1': 10, 'src2': 0}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 4 (LOAD makes b unknown, ADD not folded):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 3, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 5 (CF on generated Level 1 program, seed=42):\n", + "Before:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "Test 6 (propagation: a=0 substituted into ADD even though b unknown):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 0, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 0, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def dead_code_elimination(program):\n", + " \"\"\"\n", + " Backward liveness analysis. Replace dead instructions with NOPs,\n", + " then strip NOPs.\n", + "\n", + " Rules:\n", + " - STOREs are always live (define program output).\n", + " - Address variables used in STORE/LOAD are always live.\n", + " - Any instruction whose dest is never read later is dead.\n", + "\n", + " Returns a fresh list of instruction dicts. Never raises, never returns None.\n", + " \"\"\"\n", + " # Walk backward, build new program in reverse, then re-reverse at the end\n", + " live = set()\n", + " new_program_reversed = []\n", + "\n", + " for instr in reversed(program):\n", + " instr = instr.copy() # never alias input\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + "\n", + " if op == \"STORE\":\n", + " # STOREs are always kept; mark sources live\n", + " if isinstance(src1, str): live.add(src1)\n", + " if isinstance(dest, str): live.add(dest) # address variable\n", + " new_program_reversed.append(instr)\n", + "\n", + " elif op == \"NOP\":\n", + " # Pass through; will be stripped at the end\n", + " new_program_reversed.append(instr)\n", + "\n", + " elif op in (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\"):\n", + " if dest in live:\n", + " # Live instruction β€” keep it, mark its sources live\n", + " live.discard(dest)\n", + " if isinstance(src1, str): live.add(src1)\n", + " if isinstance(src2, str): live.add(src2)\n", + " new_program_reversed.append(instr)\n", + " else:\n", + " # Dead β€” replace with NOP\n", + " new_program_reversed.append({\n", + " \"op\": \"NOP\", \"dest\": None, \"src1\": None, \"src2\": None,\n", + " })\n", + "\n", + " else:\n", + " # Unknown op β€” defensive pass-through\n", + " new_program_reversed.append(instr)\n", + "\n", + " # Reverse back to forward order, then strip NOPs\n", + " new_program = list(reversed(new_program_reversed))\n", + " new_program = [instr for instr in new_program if instr[\"op\"] != \"NOP\"]\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "5vo-KgaffTFR" + }, + "execution_count": 12, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for dead_code_elimination ===\n", + "\n", + "# Test 1: simple dead variable\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None}, # never used\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result1 = dead_code_elimination(test1)\n", + "print(\"Test 1 (kill unused CONST):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: 'dead' instruction removed, 2 instructions remain\n", + "\n", + "# Test 2: chain of dead computation\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"x\", \"src1\": \"a\", \"src2\": \"b\"}, # x never used β†’ dead\n", + " {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 7, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result2 = dead_code_elimination(test2)\n", + "print(\"\\nTest 2 (kill unused ADD and its feeders... but only if feeders are also unused):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: x's ADD killed. a and b also killed (only fed into x, which is dead).\n", + "# Final: just CONST c=7, STORE.\n", + "\n", + "# Test 3: STORE always preserved\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result3 = dead_code_elimination(test3)\n", + "print(\"\\nTest 3 (STORE preserved, feeder kept live):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: both instructions unchanged\n", + "\n", + "# Test 4: variable used by STORE is live, even if defined far above\n", + "test4 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None}, # used by STORE β†’ live\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 9, \"src2\": None}, # never used β†’ dead\n", + " {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 1, \"src2\": None}, # never used β†’ dead\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "result4 = dead_code_elimination(test4)\n", + "print(\"\\nTest 4 (only 'a' is live, b/c killed):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: a's CONST + STORE remain. b, c stripped.\n", + "\n", + "# Test 5: combine CF + DCE on a generated program\n", + "random.seed(42)\n", + "prog = generate_level_1()\n", + "print(\"\\nTest 5 (CF then DCE on seed=42):\")\n", + "print(\"Original:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"After CF:\")\n", + "for i, instr in enumerate(after_cf): print(f\" {i}: {instr}\")\n", + "\n", + "after_dce = dead_code_elimination(after_cf)\n", + "print(\"After CF + DCE:\")\n", + "for i, instr in enumerate(after_dce): print(f\" {i}: {instr}\")\n", + "# Expected: original 4 instructions become much shorter β€” dead 'v1' eliminated, v2 folded\n", + "\n", + "# Test 6: idempotence β€” running DCE twice gives same result\n", + "test6 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n", + "]\n", + "once = dead_code_elimination(test6)\n", + "twice = dead_code_elimination(once)\n", + "print(\"\\nTest 6 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "vv8QRFRlvzdq", + "outputId": "e1d35874-ff1a-4dec-f43f-e22867dd3911" + }, + "execution_count": 8, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (kill unused CONST):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 2 (kill unused ADD and its feeders... but only if feeders are also unused):\n", + " 0: {'op': 'CONST', 'dest': 'c', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 3 (STORE preserved, feeder kept live):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 4 (only 'a' is live, b/c killed):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n", + "\n", + "Test 5 (CF then DCE on seed=42):\n", + "Original:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After CF:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After CF + DCE:\n", + " 0: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "Test 6 (idempotence): PASS\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def peephole_optimization(program):\n", + " \"\"\"\n", + " Single-instruction peephole rewrites. Replaces expensive ops with\n", + " cheaper equivalents when operands hit special values (0, 1, 2, self).\n", + "\n", + " Patterns (all preserve semantics):\n", + " MUL by 2 -> ADD x + x\n", + " MUL by 1 -> CONST (the other operand)\n", + " MUL by 0 -> CONST 0\n", + " ADD with 0 -> CONST (the other operand)\n", + " SUB x - x -> CONST 0\n", + " DIV by 1 -> CONST (the dividend)\n", + "\n", + " Notes:\n", + " - Operates on individual instructions; no cross-instruction state.\n", + " - Relies on constant_folding having propagated literal values into operands.\n", + " - Returns fresh instruction dicts; never aliases inputs.\n", + "\n", + " Returns a new list. Never raises, never returns None.\n", + " \"\"\"\n", + " new_program = []\n", + "\n", + " for instr in program:\n", + " instr = instr.copy()\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + "\n", + " # === MUL patterns ===\n", + " if op == \"MUL\":\n", + " # MUL by 0 -> CONST 0\n", + " if src1 == 0 or src2 == 0:\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n", + " continue\n", + " # MUL by 1 -> CONST (other operand) if the other operand is a literal,\n", + " # otherwise leave alone (we don't want to introduce a useless\n", + " # \"CONST dest = some_var_name\" β€” that's not a valid CONST).\n", + " if src1 == 1 and isinstance(src2, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n", + " continue\n", + " if src2 == 1 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " # MUL by 2 -> ADD x + x (when one operand is literal 2 and the other is a var)\n", + " if src1 == 2 and isinstance(src2, str):\n", + " new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src2, \"src2\": src2})\n", + " continue\n", + " if src2 == 2 and isinstance(src1, str):\n", + " new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src1, \"src2\": src1})\n", + " continue\n", + " # No pattern matched β€” keep as is\n", + " new_program.append(instr)\n", + "\n", + " # === ADD patterns ===\n", + " elif op == \"ADD\":\n", + " # ADD with 0 -> CONST (other operand) if the other operand is a literal\n", + " if src1 == 0 and isinstance(src2, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n", + " continue\n", + " if src2 == 0 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " # If \"ADD x + 0\" where x is a variable, we'd want a copy β€” but our schema\n", + " # has no MOV/COPY op. Leave alone; CF + DCE handle the rest in practice.\n", + " new_program.append(instr)\n", + "\n", + " # === SUB patterns ===\n", + " elif op == \"SUB\":\n", + " # SUB x - x -> CONST 0 (only if both sources are the same string variable)\n", + " if isinstance(src1, str) and isinstance(src2, str) and src1 == src2:\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n", + " continue\n", + " # SUB x - 0 with literal 0 on src2 β€” we'd want a copy; skip (no COPY op)\n", + " new_program.append(instr)\n", + "\n", + " # === DIV patterns ===\n", + " elif op == \"DIV\":\n", + " # DIV by 1 -> CONST (dividend) if dividend is a literal\n", + " if src2 == 1 and isinstance(src1, int):\n", + " new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n", + " continue\n", + " new_program.append(instr)\n", + "\n", + " else:\n", + " # CONST, LOAD, STORE, NOP β€” pass through\n", + " new_program.append(instr)\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "IaczDNeLv5PW" + }, + "execution_count": 9, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for peephole_optimization ===\n", + "\n", + "# Test 1: MUL by 2 -> ADD x+x\n", + "test1 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result1 = peephole_optimization(test1)\n", + "print(\"Test 1 (MUL by 2 -> ADD self):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes ADD y+y\n", + "\n", + "# Test 2: MUL by 0 -> CONST 0\n", + "test2 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 0},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result2 = peephole_optimization(test2)\n", + "print(\"\\nTest 2 (MUL by 0 -> CONST 0):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes CONST x = 0\n", + "\n", + "# Test 3: MUL by 1 (both literals)\n", + "test3 = [\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": 7, \"src2\": 1},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result3 = peephole_optimization(test3)\n", + "print(\"\\nTest 3 (MUL 7*1 -> CONST 7):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 0 becomes CONST x = 7\n", + "\n", + "# Test 4: SUB x - x -> CONST 0\n", + "test4 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"SUB\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result4 = peephole_optimization(test4)\n", + "print(\"\\nTest 4 (SUB y-y -> CONST 0):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 becomes CONST x = 0\n", + "\n", + "# Test 5: ADD with 0\n", + "test5 = [\n", + " {\"op\": \"ADD\", \"dest\": \"x\", \"src1\": 0, \"src2\": 5},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result5 = peephole_optimization(test5)\n", + "print(\"\\nTest 5 (ADD 0+5 -> CONST 5):\")\n", + "for i, instr in enumerate(result5): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 0 becomes CONST x = 5\n", + "\n", + "# Test 6: full pipeline β€” CF then peephole on a generated program\n", + "random.seed(7)\n", + "prog = generate_level_1()\n", + "print(\"\\nTest 6 (CF then peephole on seed=7):\")\n", + "print(\"Original:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"After CF:\")\n", + "for i, instr in enumerate(after_cf): print(f\" {i}: {instr}\")\n", + "after_peep = peephole_optimization(after_cf)\n", + "print(\"After CF + peephole:\")\n", + "for i, instr in enumerate(after_peep): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 7: idempotence\n", + "test7 = [\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "once = peephole_optimization(test7)\n", + "twice = peephole_optimization(once)\n", + "print(\"\\nTest 7 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")\n", + "\n", + "# Test 8: no false fires β€” vanilla program shouldn't get rewritten\n", + "test8 = [\n", + " {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"MUL\", \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"}, # y*y, no peephole pattern\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n", + "]\n", + "result8 = peephole_optimization(test8)\n", + "print(\"\\nTest 8 (no false rewrite on y*y):\")\n", + "for i, instr in enumerate(result8): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 1 unchanged (still MUL y*y)" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "DWDUZUQvwXye", + "outputId": "1d1e327f-5f26-43e4-a6ac-f7a286da3781" + }, + "execution_count": 10, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (MUL by 2 -> ADD self):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'ADD', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 2 (MUL by 0 -> CONST 0):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 3 (MUL 7*1 -> CONST 7):\n", + " 0: {'op': 'CONST', 'dest': 'x', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 4 (SUB y-y -> CONST 0):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 5 (ADD 0+5 -> CONST 5):\n", + " 0: {'op': 'CONST', 'dest': 'x', 'src1': 5, 'src2': None}\n", + " 1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n", + "\n", + "Test 6 (CF then peephole on seed=7):\n", + "Original:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "After CF:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "After CF + peephole:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "Test 7 (idempotence): PASS\n", + "\n", + "Test 8 (no false rewrite on y*y):\n", + " 0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n", + " 1: {'op': 'MUL', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n", + " 2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def dump_ir(program_data):\n", + " \"\"\"\n", + " Convert a TAC program into a readable string for the LLM observation.\n", + "\n", + " Format:\n", + " // OBSERVABLE OUT: mem[0], mem[1]\n", + " 0: v0 = 3 # CONST [1 cycle]\n", + " 1: v1 = 5 # CONST [1 cycle]\n", + " 2: v2 = v0 + v1 # ADD [1 cycle]\n", + " 3: mem[addr0] = v2 # STORE [4 cycles]\n", + " // TOTAL: 7 cycles\n", + "\n", + " Args:\n", + " program_data: dict with keys 'instructions' and 'observable_addrs',\n", + " OR a raw list of instructions (legacy).\n", + "\n", + " Returns:\n", + " A multi-line string suitable for inclusion in an LLM prompt.\n", + " \"\"\"\n", + " # Accept both forms β€” full program dict or raw instruction list\n", + " if isinstance(program_data, dict):\n", + " instructions = program_data[\"instructions\"]\n", + " observable_addrs = program_data.get(\"observable_addrs\", [])\n", + " else:\n", + " instructions = program_data\n", + " observable_addrs = []\n", + "\n", + " lines = []\n", + "\n", + " # Header: observable outputs\n", + " if observable_addrs:\n", + " addr_str = \", \".join(f\"mem[{a}]\" for a in observable_addrs)\n", + " lines.append(f\"// OBSERVABLE OUT: {addr_str}\")\n", + "\n", + " # Body: each instruction in human-readable form\n", + " total_cycles = 0\n", + " for i, instr in enumerate(instructions):\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + " src1 = instr[\"src1\"]\n", + " src2 = instr[\"src2\"]\n", + " cost = CYCLE_COSTS.get(op, 0)\n", + " total_cycles += cost\n", + "\n", + " # Render the instruction body\n", + " if op == \"CONST\":\n", + " body = f\"{dest} = {src1}\"\n", + " elif op == \"ADD\":\n", + " body = f\"{dest} = {src1} + {src2}\"\n", + " elif op == \"SUB\":\n", + " body = f\"{dest} = {src1} - {src2}\"\n", + " elif op == \"MUL\":\n", + " body = f\"{dest} = {src1} * {src2}\"\n", + " elif op == \"DIV\":\n", + " body = f\"{dest} = {src1} // {src2}\"\n", + " elif op == \"LOAD\":\n", + " body = f\"{dest} = mem[{src1}]\"\n", + " elif op == \"STORE\":\n", + " body = f\"mem[{dest}] = {src1}\"\n", + " elif op == \"NOP\":\n", + " body = \"nop\"\n", + " else:\n", + " body = f\"\"\n", + "\n", + " cost_label = f\"{cost} cycle\" if cost == 1 else f\"{cost} cycles\"\n", + " lines.append(f\"{i}: {body:<35} # {op:<6} [{cost_label}]\")\n", + "\n", + " # Footer: total cycles\n", + " lines.append(f\"// TOTAL: {total_cycles} cycles\")\n", + "\n", + " return \"\\n\".join(lines)" + ], + "metadata": { + "id": "yB_yMNOZwaPT" + }, + "execution_count": 13, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for state translator ===\n", + "\n", + "# Test 1: full generated program\n", + "random.seed(1)\n", + "prog = generate_level_1()\n", + "print(\"Test 1 (raw seed=1 program):\")\n", + "print(dump_ir(prog))\n", + "\n", + "# Test 2: after CF\n", + "print(\"\\nTest 2 (after CF):\")\n", + "folded = constant_folding(prog[\"instructions\"])\n", + "prog_after_cf = {**prog, \"instructions\": folded}\n", + "print(dump_ir(prog_after_cf))\n", + "\n", + "# Test 3: after CF + DCE β€” should show fewer instructions, lower total\n", + "print(\"\\nTest 3 (after CF + DCE):\")\n", + "optimized = dead_code_elimination(folded)\n", + "prog_optimized = {**prog, \"instructions\": optimized}\n", + "print(dump_ir(prog_optimized))\n", + "\n", + "# Test 4: every op type at least once\n", + "test4_instructions = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"b\", \"src1\": \"a\", \"src2\": 3},\n", + " {\"op\": \"SUB\", \"dest\": \"c\", \"src1\": \"b\", \"src2\": \"a\"},\n", + " {\"op\": \"MUL\", \"dest\": \"d\", \"src1\": \"c\", \"src2\": 2},\n", + " {\"op\": \"DIV\", \"dest\": \"e\", \"src1\": \"d\", \"src2\": 4},\n", + " {\"op\": \"LOAD\", \"dest\": \"f\", \"src1\": \"addr0\", \"src2\": None},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"e\", \"src2\": None},\n", + " {\"op\": \"NOP\", \"dest\": None, \"src1\": None, \"src2\": None},\n", + "]\n", + "test4_data = {\n", + " \"instructions\": test4_instructions,\n", + " \"observable_addrs\": [0, 1],\n", + "}\n", + "print(\"\\nTest 4 (every op type):\")\n", + "print(dump_ir(test4_data))\n", + "\n", + "# Test 5: empty program shouldn't crash\n", + "print(\"\\nTest 5 (empty program):\")\n", + "print(dump_ir({\"instructions\": [], \"observable_addrs\": [0]}))" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "SUqa7odp3wB_", + "outputId": "c3f44089-213a-4a69-f29f-42b755ed142a" + }, + "execution_count": 14, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (raw seed=1 program):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 2 # CONST [1 cycle]\n", + "2: v2 = v1 + v1 # ADD [1 cycle]\n", + "3: v3 = v2 * v1 # MUL [3 cycles]\n", + "4: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 10 cycles\n", + "\n", + "Test 2 (after CF):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v1 = 2 # CONST [1 cycle]\n", + "2: v2 = 4 # CONST [1 cycle]\n", + "3: v3 = 8 # CONST [1 cycle]\n", + "4: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 8 cycles\n", + "\n", + "Test 3 (after CF + DCE):\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v3 = 8 # CONST [1 cycle]\n", + "1: mem[addr0] = v3 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n", + "\n", + "Test 4 (every op type):\n", + "// OBSERVABLE OUT: mem[0], mem[1]\n", + "0: a = 5 # CONST [1 cycle]\n", + "1: b = a + 3 # ADD [1 cycle]\n", + "2: c = b - a # SUB [1 cycle]\n", + "3: d = c * 2 # MUL [3 cycles]\n", + "4: e = d // 4 # DIV [5 cycles]\n", + "5: f = mem[addr0] # LOAD [4 cycles]\n", + "6: mem[addr0] = e # STORE [4 cycles]\n", + "7: nop # NOP [0 cycles]\n", + "// TOTAL: 19 cycles\n", + "\n", + "Test 5 (empty program):\n", + "// OBSERVABLE OUT: mem[0]\n", + "// TOTAL: 0 cycles\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def generate_level_2():\n", + " \"\"\"\n", + " Generate a Level 2 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 8-12 instructions\n", + " - 3-5 CONSTs (some literal, some used in arithmetic)\n", + " - 3-5 arithmetic ops (ADD, SUB, MUL, DIV) with chaining\n", + " - 1-3 dead variables (DCE opportunities)\n", + " - 1 LOAD from initial memory (introduces non-constant variable)\n", + " - 1 STORE at the end (observable output)\n", + " - Mix of operands designed to hit peephole patterns occasionally:\n", + " - MUL by literal 2 / 1 / 0 (not always β€” diversity matters)\n", + " - DIV by non-zero divisor only (Design Assumption #4)\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " available_vars = [] # vars that exist and can be used as sources\n", + "\n", + " # Step 1: 3-5 CONSTs\n", + " num_consts = random.randint(3, 5)\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " available_vars.append(var)\n", + "\n", + " # Step 2: 1 LOAD from initial memory\n", + " # We'll seed initial_mem with something at a fresh address.\n", + " load_addr_var = \"addr_in\"\n", + " loaded_var = new_var()\n", + " instructions.append({\n", + " \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": load_addr_var, \"src2\": None,\n", + " })\n", + " available_vars.append(loaded_var)\n", + "\n", + " # Step 3: 3-5 arithmetic ops, chained\n", + " num_arith = random.randint(3, 5)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n", + "\n", + " # 30% chance to use a literal as src2 to expose peephole opportunities\n", + " # (MUL by 2/1/0, etc.)\n", + " use_literal_src2 = random.random() < 0.3\n", + "\n", + " src1 = random.choice(available_vars)\n", + " if use_literal_src2:\n", + " # For DIV, never emit literal 0 (Design Assumption #4)\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4]) # safe non-zero divisors\n", + " else:\n", + " src2 = random.choice([0, 1, 2, 3]) # 0/1/2 hit peephole patterns\n", + " else:\n", + " src2 = random.choice(available_vars)\n", + " # If we sampled a variable for DIV's src2, we can't be sure it's non-zero\n", + " # at runtime. To stay safe per Assumption #4, restrict DIV to literal src2.\n", + " if op == \"DIV\":\n", + " src2 = random.choice([1, 2, 3, 4])\n", + "\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " available_vars.append(dest)\n", + " last_result = dest\n", + "\n", + " # Step 4: 1-3 extra dead CONSTs sprinkled in (DCE targets)\n", + " num_dead = random.randint(1, 3)\n", + " for _ in range(num_dead):\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 20)\n", + " # Insert at a random position before the (eventual) STORE\n", + " insert_pos = random.randint(0, len(instructions))\n", + " instructions.insert(insert_pos, {\n", + " \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n", + " })\n", + " # Note: dead_var intentionally never used afterward\n", + "\n", + " # Step 5: STORE the final result\n", + " initial_vars = {\n", + " \"addr0\": 0,\n", + " \"addr_in\": 1, # address that the LOAD reads from\n", + " }\n", + " initial_mem = {1: random.randint(1, 20)} # seed mem[1] with a random value\n", + "\n", + " instructions.append({\n", + " \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n", + " })\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": [0],\n", + " }" + ], + "metadata": { + "id": "OL6QbL0H3xuP" + }, + "execution_count": 15, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for Level 2 generator ===\n", + "\n", + "# Test 1: spot-check a few seeds\n", + "for seed in [42, 1, 7]:\n", + " random.seed(seed)\n", + " prog = generate_level_2()\n", + " print(f\"\\n=== Level 2, seed={seed} ===\")\n", + " print(dump_ir(prog))\n", + " print(f\"initial_vars: {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + "\n", + "# Test 2: full pipeline (CF + DCE + peephole) on a Level 2 program\n", + "print(\"\\n=== Full optimization pipeline on Level 2 (seed=42) ===\")\n", + "random.seed(42)\n", + "prog = generate_level_2()\n", + "print(\"ORIGINAL:\")\n", + "print(dump_ir(prog))\n", + "\n", + "after_cf = constant_folding(prog[\"instructions\"])\n", + "print(\"\\nAFTER CF:\")\n", + "print(dump_ir({**prog, \"instructions\": after_cf}))\n", + "\n", + "after_dce = dead_code_elimination(after_cf)\n", + "print(\"\\nAFTER CF + DCE:\")\n", + "print(dump_ir({**prog, \"instructions\": after_dce}))\n", + "\n", + "after_peep = peephole_optimization(after_dce)\n", + "print(\"\\nAFTER CF + DCE + PEEPHOLE:\")\n", + "print(dump_ir({**prog, \"instructions\": after_peep}))" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "D0XjNetU5gIo", + "outputId": "1c35cdfc-a95d-48ff-c7ce-c4c236bbbbdb" + }, + "execution_count": 16, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== Level 2, seed=42 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 23 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 1}\n", + "\n", + "=== Level 2, seed=1 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 10 # CONST [1 cycle]\n", + "1: v7 = 19 # CONST [1 cycle]\n", + "2: v1 = 2 # CONST [1 cycle]\n", + "3: v2 = 5 # CONST [1 cycle]\n", + "4: v3 = mem[addr_in] # LOAD [4 cycles]\n", + "5: v4 = v3 // 2 # DIV [5 cycles]\n", + "6: v5 = v3 + v3 # ADD [1 cycle]\n", + "7: v6 = v2 + v5 # ADD [1 cycle]\n", + "8: mem[addr0] = v6 # STORE [4 cycles]\n", + "// TOTAL: 19 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 11}\n", + "\n", + "=== Level 2, seed=7 ===\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 3 # CONST [1 cycle]\n", + "1: v1 = 7 # CONST [1 cycle]\n", + "2: v2 = 1 # CONST [1 cycle]\n", + "3: v3 = 2 # CONST [1 cycle]\n", + "4: v4 = mem[addr_in] # LOAD [4 cycles]\n", + "5: v5 = v0 + v4 # ADD [1 cycle]\n", + "6: v6 = v3 - 3 # SUB [1 cycle]\n", + "7: v7 = v4 + 3 # ADD [1 cycle]\n", + "8: v10 = 2 # CONST [1 cycle]\n", + "9: v8 = v1 + v3 # ADD [1 cycle]\n", + "10: v9 = v6 + v0 # ADD [1 cycle]\n", + "11: mem[addr0] = v9 # STORE [4 cycles]\n", + "// TOTAL: 18 cycles\n", + "initial_vars: {'addr0': 0, 'addr_in': 1}\n", + "initial_mem : {1: 5}\n", + "\n", + "=== Full optimization pipeline on Level 2 (seed=42) ===\n", + "ORIGINAL:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = v4 + v0 # ADD [1 cycle]\n", + "10: v7 = v0 // 2 # DIV [5 cycles]\n", + "11: v8 = v0 - v3 # SUB [1 cycle]\n", + "12: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 23 cycles\n", + "\n", + "AFTER CF:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v0 = 2 # CONST [1 cycle]\n", + "1: v1 = 1 # CONST [1 cycle]\n", + "2: v2 = 5 # CONST [1 cycle]\n", + "3: v3 = 4 # CONST [1 cycle]\n", + "4: v11 = 19 # CONST [1 cycle]\n", + "5: v4 = 4 # CONST [1 cycle]\n", + "6: v5 = mem[addr_in] # LOAD [4 cycles]\n", + "7: v9 = 18 # CONST [1 cycle]\n", + "8: v10 = 8 # CONST [1 cycle]\n", + "9: v6 = 6 # CONST [1 cycle]\n", + "10: v7 = 1 # CONST [1 cycle]\n", + "11: v8 = -2 # CONST [1 cycle]\n", + "12: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 19 cycles\n", + "\n", + "AFTER CF + DCE:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v8 = -2 # CONST [1 cycle]\n", + "1: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n", + "\n", + "AFTER CF + DCE + PEEPHOLE:\n", + "// OBSERVABLE OUT: mem[0]\n", + "0: v8 = -2 # CONST [1 cycle]\n", + "1: mem[addr0] = v8 # STORE [4 cycles]\n", + "// TOTAL: 5 cycles\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [], + "metadata": { + "id": "Vo-9MitD5hwM" + }, + "execution_count": null, + "outputs": [] + } + ] +} \ No newline at end of file diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1.ipynb b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..448c084a417e1fb58a2a89826df24cf2f96e616c --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1.ipynb @@ -0,0 +1,484 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "code", + "source": [ + "## Design Assumptions (do not violate)\n", + "\n", + "# 1. **DCE never eliminates STOREs.** They define program output (final mem state).\n", + "# 2. **Addresses are distinct by construction.** Generator allocates each address variable to a unique integer; no aliasing.\n", + "# 3. **CF refuses to fold DIV by zero.** If src2 == 0 on a DIV op, leave instruction unchanged.\n", + "# 4. **Generator never emits DIV by literal zero.** When DIV is generated, src2 is always a non-zero constant or a variable known to be non-zero.\n", + "# # 5. **Integer arithmetic only.** No floats anywhere β€” avoids equivalence-check precision issues.\n", + "# 6. Generator declares `observable_addrs` per program β€” verifier compares only these mem entries.\n", + "# 7. State translator annotates observable outputs at top of dump.\n", + "# 8. Reward distinguishes broken (-1000) from valid-but-worse (small negative) β€” Harshal's formula.\n", + "# 9. Multi-input verification: 3-5 random initial states, all must match.\n", + "# 10. Integer division uses Python floor division (//). Aarush's VM must match." + ], + "metadata": { + "id": "_XI5jT2Ibvrf" + }, + "execution_count": 17, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "id": "j6CQ327KWXwr" + }, + "outputs": [], + "source": [ + "# === TAC Schema v1.0 (LOCKED with Role 1 / Aarush) ===\n", + "# Reverse passes deferred to stretch goal β€” not in initial action space.\n", + "\n", + "OPS = [\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\", \"STORE\", \"NOP\"]\n", + "\n", + "CYCLE_COSTS = {\n", + " \"CONST\": 1,\n", + " \"ADD\": 1,\n", + " \"SUB\": 1,\n", + " \"MUL\": 3,\n", + " \"DIV\": 5,\n", + " \"LOAD\": 4,\n", + " \"STORE\": 4,\n", + " \"NOP\": 0,\n", + "}\n", + "\n", + "# Instruction shape: {\"op\": str, \"dest\": str|None, \"src1\": Any, \"src2\": Any}\n", + "# Operands: str = variable name, int = literal constant, None = unused\n", + "#\n", + "# Op semantics:\n", + "# CONST: dest = src1 (src1 is int literal, src2 = None)\n", + "# ADD/SUB/MUL/DIV: dest = src1 OP src2 (src1, src2 are var names or int literals)\n", + "# LOAD: dest = mem[src1] (src1 is a var holding an address)\n", + "# STORE: mem[dest] = src1 (dest is a var holding an address)\n", + "# NOP: no-op (all fields None)\n", + "#\n", + "# Program output (for equivalence check) = final memory state (mem dict).\n", + "# Programs ship as: (initial_vars: dict, initial_mem: dict, instructions: list[dict])\n", + "# 10. Integer division uses Python floor division (//). Aarush's VM must match." + ] + }, + { + "cell_type": "code", + "source": [ + "import random\n", + "\n", + "def generate_level_1():\n", + " \"\"\"\n", + " Generate a Level 1 Toy-IR program.\n", + "\n", + " Characteristics:\n", + " - 4-6 instructions\n", + " - 2-3 CONST ops with literal values\n", + " - 1-2 arithmetic ops on those constants (foldable by CF)\n", + " - 0-1 dead variables (killable by DCE)\n", + " - Exactly 1 STORE at the end so the program has an observable output\n", + "\n", + " Returns:\n", + " dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n", + " \"\"\"\n", + " instructions = []\n", + " var_counter = 0\n", + "\n", + " def new_var():\n", + " nonlocal var_counter\n", + " name = f\"v{var_counter}\"\n", + " var_counter += 1\n", + " return name\n", + "\n", + " # Step 1: Generate 2-3 constant assignments\n", + " num_consts = random.randint(2, 3)\n", + " const_vars = []\n", + " for _ in range(num_consts):\n", + " var = new_var()\n", + " value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n", + " const_vars.append(var)\n", + "\n", + " # Step 2: Generate 1-2 arithmetic ops using those constants\n", + " num_arith = random.randint(1, 2)\n", + " last_result = None\n", + " for _ in range(num_arith):\n", + " op = random.choice([\"ADD\", \"MUL\"])\n", + " src1 = random.choice(const_vars)\n", + " src2 = random.choice(const_vars)\n", + " dest = new_var()\n", + " instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n", + " last_result = dest\n", + " const_vars.append(dest)\n", + "\n", + " # Step 3: Optionally add 1 dead variable (50% chance)\n", + " if random.random() < 0.5:\n", + " dead_var = new_var()\n", + " dead_value = random.randint(1, 10)\n", + " instructions.append({\"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None})\n", + " # Note: dead_var is intentionally never used β€” DCE should catch it.\n", + "\n", + " # Step 4: Add a STORE at the end so the program has observable output\n", + " initial_vars = {\"addr0\": 0}\n", + " instructions.append({\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None})\n", + "\n", + " initial_mem = {}\n", + "\n", + " return {\n", + " \"initial_vars\": initial_vars,\n", + " \"initial_mem\": initial_mem,\n", + " \"instructions\": instructions,\n", + " \"observable_addrs\": [0], # Aarush's verifier compares only these mem entries\n", + " }\n", + "\n", + "\n", + "# Sanity-check: generate a few programs and print them\n", + "for seed in [42, 1, 7, 99]:\n", + " random.seed(seed)\n", + " prog = generate_level_1()\n", + " print(f\"\\n=== seed={seed} ===\")\n", + " print(f\"initial_vars : {prog['initial_vars']}\")\n", + " print(f\"initial_mem : {prog['initial_mem']}\")\n", + " print(f\"observable_addrs : {prog['observable_addrs']}\")\n", + " print(f\"instructions ({len(prog['instructions'])}):\")\n", + " for i, instr in enumerate(prog['instructions']):\n", + " print(f\" {i}: {instr}\")" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "a86vD1OyWcB9", + "outputId": "de3d7093-eaea-420e-a0d7-2bd9451bcf6b" + }, + "execution_count": 14, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "=== seed=42 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (4):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "=== seed=1 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 2, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v1', 'src2': 'v1'}\n", + " 3: {'op': 'MUL', 'dest': 'v3', 'src1': 'v2', 'src2': 'v1'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=7 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (6):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n", + " 4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n", + " 5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n", + "\n", + "=== seed=99 ===\n", + "initial_vars : {'addr0': 0}\n", + "initial_mem : {}\n", + "observable_addrs : [0]\n", + "instructions (5):\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 7, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 4, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 10, 'src2': None}\n", + " 3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v0', 'src2': 'v0'}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [ + "def _resolve_operand(operand, known_constants):\n", + " \"\"\"\n", + " Given a TAC operand (string variable name or int literal),\n", + " return its concrete integer value if known, or None if unknown.\n", + " \"\"\"\n", + " if isinstance(operand, int):\n", + " return operand\n", + " if isinstance(operand, str) and operand in known_constants:\n", + " return known_constants[operand]\n", + " return None\n", + "\n", + "\n", + "def _compute(op, v1, v2):\n", + " \"\"\"Compute the result of a binary arithmetic op on two known integers.\"\"\"\n", + " if op == \"ADD\": return v1 + v2\n", + " if op == \"SUB\": return v1 - v2\n", + " if op == \"MUL\": return v1 * v2\n", + " if op == \"DIV\": return v1 // v2 # floor division (locked semantics)\n", + " raise ValueError(f\"_compute called with non-arithmetic op: {op}\")\n", + "\n", + "\n", + "def constant_folding(program):\n", + " \"\"\"\n", + " Forward pass that folds constant arithmetic into CONST ops, and\n", + " propagates known constants into instruction operands.\n", + "\n", + " Behavior:\n", + " - If both sources of an arithmetic op resolve to known integers,\n", + " replaces the instruction with a CONST holding the computed result.\n", + " - If only one source is known, still substitutes that known value\n", + " into the instruction (constant propagation), enabling downstream\n", + " passes (e.g., peephole) to recognize patterns like ADD-with-0 or MUL-by-1.\n", + " - Refuses to fold DIV by zero (Design Assumption #3).\n", + " - Always returns fresh dicts; never aliases input instructions.\n", + "\n", + " Args:\n", + " program: list of TAC instruction dicts (per locked schema).\n", + "\n", + " Returns:\n", + " new list of TAC instruction dicts. Always semantics-preserving.\n", + " Never raises, never returns None.\n", + " \"\"\"\n", + " known_constants = {}\n", + " new_program = []\n", + "\n", + " for instr in program:\n", + " # Always work on a copy β€” never alias input dicts\n", + " instr = instr.copy()\n", + " op = instr[\"op\"]\n", + " dest = instr[\"dest\"]\n", + "\n", + " if op == \"CONST\":\n", + " known_constants[dest] = instr[\"src1\"]\n", + " new_program.append(instr)\n", + "\n", + " elif op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n", + " # === Constant propagation: substitute known constants into operands ===\n", + " if isinstance(instr[\"src1\"], str) and instr[\"src1\"] in known_constants:\n", + " instr[\"src1\"] = known_constants[instr[\"src1\"]]\n", + " if isinstance(instr[\"src2\"], str) and instr[\"src2\"] in known_constants:\n", + " instr[\"src2\"] = known_constants[instr[\"src2\"]]\n", + "\n", + " # === Try to fold ===\n", + " v1 = _resolve_operand(instr[\"src1\"], known_constants)\n", + " v2 = _resolve_operand(instr[\"src2\"], known_constants)\n", + "\n", + " if v1 is not None and v2 is not None:\n", + " # Both operands are known integers\n", + " if op == \"DIV\" and v2 == 0:\n", + " # Refuse to fold DIV by zero\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + " else:\n", + " # Fold: replace with CONST\n", + " result = _compute(op, v1, v2)\n", + " new_program.append({\n", + " \"op\": \"CONST\",\n", + " \"dest\": dest,\n", + " \"src1\": result,\n", + " \"src2\": None,\n", + " })\n", + " known_constants[dest] = result\n", + " else:\n", + " # Can't fold (at least one operand unknown).\n", + " # Instruction may still have been mutated by propagation above.\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op == \"LOAD\":\n", + " # Memory reads aren't statically resolvable\n", + " new_program.append(instr)\n", + " known_constants.pop(dest, None)\n", + "\n", + " elif op in (\"STORE\", \"NOP\"):\n", + " # No dest tracking needed.\n", + " # Note: we COULD propagate src1 into a STORE for downstream readability,\n", + " # but the cycle cost is unchanged and the executor handles vars fine.\n", + " # Leave STORE alone β€” keeps the code minimal.\n", + " new_program.append(instr)\n", + "\n", + " else:\n", + " # Unknown op β€” defensive pass-through\n", + " new_program.append(instr)\n", + "\n", + " return new_program" + ], + "metadata": { + "id": "CGk8Fz4CZA3t" + }, + "execution_count": 18, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# === Sanity tests for constant_folding ===\n", + "\n", + "# Test 1: simple fold β€” ADD of two CONSTs\n", + "test1 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result1 = constant_folding(test1)\n", + "print(\"Test 1 (simple ADD fold):\")\n", + "for i, instr in enumerate(result1): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 becomes CONST c = 8\n", + "\n", + "# Test 2: chained fold β€” second op uses first op's folded result\n", + "test2 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # β†’ 8\n", + " {\"op\": \"MUL\", \"dest\": \"d\", \"src1\": \"c\", \"src2\": \"b\"}, # β†’ 8 * 5 = 40\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"d\", \"src2\": None},\n", + "]\n", + "result2 = constant_folding(test2)\n", + "print(\"\\nTest 2 (chained fold):\")\n", + "for i, instr in enumerate(result2): print(f\" {i}: {instr}\")\n", + "# Expected: c becomes CONST 8, d becomes CONST 40\n", + "\n", + "# Test 3: DIV by zero refused\n", + "test3 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 10, \"src2\": None},\n", + " {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"DIV\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # 10/0 β€” must NOT fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result3 = constant_folding(test3)\n", + "print(\"\\nTest 3 (DIV by zero refused):\")\n", + "for i, instr in enumerate(result3): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still DIV, not CONST)\n", + "\n", + "# Test 4: unknown source can't fold\n", + "test4 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b is unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # can't fold\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result4 = constant_folding(test4)\n", + "print(\"\\nTest 4 (LOAD makes b unknown, ADD not folded):\")\n", + "for i, instr in enumerate(result4): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 unchanged (still ADD)\n", + "\n", + "# Test 5: idempotence β€” running CF on a generated program\n", + "random.seed(42)\n", + "prog = generate_level_1()\n", + "folded = constant_folding(prog[\"instructions\"])\n", + "print(\"\\nTest 5 (CF on generated Level 1 program, seed=42):\")\n", + "print(\"Before:\")\n", + "for i, instr in enumerate(prog[\"instructions\"]): print(f\" {i}: {instr}\")\n", + "print(\"After:\")\n", + "for i, instr in enumerate(folded): print(f\" {i}: {instr}\")\n", + "\n", + "# Test 6: constant propagation β€” only one source is known\n", + "test6 = [\n", + " {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 0, \"src2\": None},\n", + " {\"op\": \"LOAD\", \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None}, # b unknown\n", + " {\"op\": \"ADD\", \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"}, # a known (=0), b unknown\n", + " {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n", + "]\n", + "result6 = constant_folding(test6)\n", + "print(\"\\nTest 6 (propagation: a=0 substituted into ADD even though b unknown):\")\n", + "for i, instr in enumerate(result6): print(f\" {i}: {instr}\")\n", + "# Expected: instruction 2 is still ADD (can't fold β€” b unknown), but src1 is now literal 0, not 'a'\n", + "# This sets up peephole to recognize \"ADD with 0\" later" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "WLKSTVa8fRF8", + "outputId": "27042d38-0044-4814-bbfa-8a2c2d87ee06" + }, + "execution_count": 20, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Test 1 (simple ADD fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 2 (chained fold):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n", + " 3: {'op': 'CONST', 'dest': 'd', 'src1': 40, 'src2': None}\n", + " 4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'd', 'src2': None}\n", + "\n", + "Test 3 (DIV by zero refused):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 10, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'b', 'src1': 0, 'src2': None}\n", + " 2: {'op': 'DIV', 'dest': 'c', 'src1': 10, 'src2': 0}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 4 (LOAD makes b unknown, ADD not folded):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 3, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n", + "\n", + "Test 5 (CF on generated Level 1 program, seed=42):\n", + "Before:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "After:\n", + " 0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n", + " 1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n", + " 2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n", + "\n", + "Test 6 (propagation: a=0 substituted into ADD even though b unknown):\n", + " 0: {'op': 'CONST', 'dest': 'a', 'src1': 0, 'src2': None}\n", + " 1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n", + " 2: {'op': 'ADD', 'dest': 'c', 'src1': 0, 'src2': 'b'}\n", + " 3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n" + ] + } + ] + }, + { + "cell_type": "code", + "source": [], + "metadata": { + "id": "5vo-KgaffTFR" + }, + "execution_count": null, + "outputs": [] + } + ] +} \ No newline at end of file diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.github/workflows/deploy.yml b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.github/workflows/deploy.yml new file mode 100644 index 0000000000000000000000000000000000000000..431f6102f3b8d75ff8565742d900027a20ac9fe1 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.github/workflows/deploy.yml @@ -0,0 +1,30 @@ +name: Deploy to Hugging Face Space + +on: + push: + branches: + - main + +jobs: + deploy: + runs-on: ubuntu-latest + + steps: + - name: Checkout Repo + uses: actions/checkout@v4 + + - name: Push to Hugging Face + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + git config --global user.email "github-actions@github.com" + git config --global user.name "github-actions" + + git clone https://user:$HF_TOKEN@huggingface.co/spaces/greedybeserk95/Compilertetris space + + rsync -av --exclude ".git" ./ space/ + + cd space + git add . + git commit -m "Auto deploy from GitHub" || echo "No changes" + git push diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md index b56fb9d2efcaabd99b42917ff3c7dc0dedab8d38..991a6a7bfcba7747510a50df55f2a62e2281cb1c 100644 --- a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md @@ -1,11 +1 @@ ---- -title: Compilertetris -emoji: ⚑ -colorFrom: red -colorTo: gray -sdk: docker -pinned: false -short_description: an RL environment which optimizes IR code ---- - -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# MetaHackathon2026Finals \ No newline at end of file diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.gitattributes b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b56fb9d2efcaabd99b42917ff3c7dc0dedab8d38 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md @@ -0,0 +1,11 @@ +--- +title: Compilertetris +emoji: ⚑ +colorFrom: red +colorTo: gray +sdk: docker +pinned: false +short_description: an RL environment which optimizes IR code +--- + +Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference diff --git a/space/space/space/space/space/space/space/space/space/space/train.py b/space/space/space/space/space/space/space/space/space/space/train.py new file mode 100644 index 0000000000000000000000000000000000000000..dd5156e8b3b0c703f0a145dd18d15ebe02876b23 --- /dev/null +++ b/space/space/space/space/space/space/space/space/space/space/train.py @@ -0,0 +1,208 @@ +""" +Toy REINFORCE training for CompilerOptimizationEnv (CPU, no PyTorch). + +This is meant for the Hugging Face Space: fast smoke training that uses the +same `runtime_core` mock engine and passes as the Gradio demo. + +For full GRPO + Unsloth + LLM, use your Colab / GPU notebooks +(`compiler_optimization_grpo.ipynb`, `role2_deliverable3_*`). + +CLI: python train.py --episodes 50 --max-steps 8 --seed 0 +Or import: from train import run_toy_training; print(run_toy_training(20, 8, 0)) +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import textwrap +from typing import List, Sequence, Tuple + +from runtime_core import ( + CompilerOptimizationEnv, + MOCK_PASSES, + MockEngine, + SAMPLE_PROGRAM, +) + +# Additional tiny programs so the policy is not overfit to a single IR. +TRAINING_PROGRAMS: List[List[dict]] = [ + SAMPLE_PROGRAM, + [ + {"op": "const", "dest": "a", "args": ["1"], "type": "int"}, + {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"}, + {"op": "ret", "args": ["b"]}, + ], + [ + {"op": "const", "dest": "x", "args": ["2"], "type": "int"}, + {"op": "const", "dest": "y", "args": ["4"], "type": "int"}, + {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"}, + {"op": "const", "dest": "k", "args": ["1"], "type": "int"}, + {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"}, + {"op": "ret", "args": ["w"]}, + ], +] + +ACTIONS: Tuple[str, ...] = tuple(sorted(MOCK_PASSES.keys())) + + +def _softmax(logits: Sequence[float]) -> List[float]: + m = max(logits) if logits else 0.0 + ex = [math.exp(x - m) for x in logits] + s = sum(ex) or 1.0 + return [e / s for e in ex] + + +def _sample_action(rng: random.Random, logits: List[float]) -> Tuple[int, List[float], float]: + """Return (action index, prob vector, log prob of chosen action).""" + p = _softmax(logits) + u = rng.random() + acc = 0.0 + idx = len(p) - 1 + for i, pi in enumerate(p): + acc += pi + if u <= acc: + idx = i + break + log_p = math.log(p[idx] + 1e-12) + return idx, p, log_p + + +def _rollout( + engine: MockEngine, + program: List[dict], + max_steps: int, + logits: List[float], + rng: random.Random, +) -> Tuple[float, List[Tuple[int, List[float], float]]]: + """ + One episode. Returns total reward and per-step (action index, prob vector, step reward) + for REINFORCE with returns G_t = sum of rewards from t onward. + """ + env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=max_steps) + env.reset(program) + total = 0.0 + trace: List[Tuple[int, List[float], float]] = [] + for _ in range(max_steps): + a_idx, pvec, _ = _sample_action(rng, logits) + action = ACTIONS[a_idx] + step = env.step(action) + r = float(step.reward) + total += r + trace.append((a_idx, pvec, r)) + if step.done: + break + return total, trace + + +def _reinforce_update( + logits: List[float], trace: List[Tuple[int, List[float], float]], lr: float +) -> None: + """REINFORCE with Monte Carlo return G_t from each step.""" + G = 0.0 + for t in range(len(trace) - 1, -1, -1): + a_idx, pvec, r = trace[t] + G += r + for i in range(len(logits)): + delta = 1.0 if i == a_idx else 0.0 + logits[i] += lr * G * (delta - pvec[i]) + + +def run_toy_training( + episodes: int, + max_steps: int, + seed: int = 0, + lr: float = 0.15, +) -> str: + """ + Train a stateless categorical policy over the three mock passes; print-friendly report. + """ + if episodes < 1: + return "episodes must be >= 1" + if max_steps < 1: + return "max_steps must be >= 1" + episodes = int(episodes) + max_steps = int(max_steps) + seed = int(seed) + lr = float(lr) + + rng = random.Random(seed) + engine = MockEngine() + logits = [0.0 for _ in ACTIONS] + + history: List[Tuple[int, float]] = [] + for ep in range(1, episodes + 1): + program = TRAINING_PROGRAMS[(ep - 1) % len(TRAINING_PROGRAMS)] + G, trace = _rollout(engine, program, max_steps, logits, rng) + if trace: + _reinforce_update(logits, trace, lr) + history.append((ep, G)) + + final_p = _softmax(logits) + lines = [ + "Toy REINFORCE (stateless policy over pass names, CPU, stdlib only)", + f" episodes={episodes} max_steps={max_steps} seed={seed} lr={lr}", + f" actions order: {list(ACTIONS)}", + "", + f" final logits: {[round(x, 4) for x in logits]}", + f" final policy: {', '.join(f'{a}={p:.3f}' for a, p in zip(ACTIONS, final_p))}", + "", + " return per episode (last 10): " + ", ".join(f"{G:+.1f}" for _, G in history[-10:]), + "", + ] + if history: + mean_r = sum(G for _, G in history) / len(history) + lines.append(f" mean return over all episodes: {mean_r:+.3f}") + lines.append("") + lines.append( + textwrap.dedent( + """ + This does not train an LLM. For GRPO + Qwen + Unsloth, run the project notebooks + on a GPU machine (e.g. Colab), not the CPU Space. + """ + ).strip() + ) + return "\n".join(lines) + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Toy REINFORCE on CompilerOptimizationEnv (CPU)") + p.add_argument("--episodes", type=int, default=50, help="Number of training episodes") + p.add_argument("--max-steps", type=int, default=8, help="max_steps per env episode") + p.add_argument("--seed", type=int, default=0, help="RNG seed") + p.add_argument("--lr", type=float, default=0.15, help="REINFORCE learning rate") + p.add_argument( + "--out-json", + type=str, + default="", + help="If set, write a small run summary to this path (e.g. training_log.json).", + ) + return p.parse_args() + + +def main() -> None: + args = _parse_args() + report = run_toy_training( + episodes=args.episodes, + max_steps=args.max_steps, + seed=args.seed, + lr=args.lr, + ) + print(report) + if args.out_json: + payload = { + "episodes": args.episodes, + "max_steps": args.max_steps, + "seed": args.seed, + "lr": args.lr, + "summary_text": report, + } + with open(args.out_json, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2) + print(f"\nWrote {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/space/space/space/space/space/space/space/space/untitled folder.zip b/space/space/space/space/space/space/space/space/untitled folder.zip new file mode 100644 index 0000000000000000000000000000000000000000..5a442bd5824eb85eab391d129ee5bbce8e6a5eed --- /dev/null +++ b/space/space/space/space/space/space/space/space/untitled folder.zip @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b29f9a98dd93fb2a14b1d3e8409ce51a756b578debde97af1190991bbb65d203 +size 63307 diff --git a/space/space/space/space/space/space/write_colab_synth_notebook.py b/space/space/space/space/space/space/write_colab_synth_notebook.py new file mode 100644 index 0000000000000000000000000000000000000000..5149477f96f750c894d3f896a70288afb9103b15 --- /dev/null +++ b/space/space/space/space/space/space/write_colab_synth_notebook.py @@ -0,0 +1,345 @@ +"""Generate Compilertetris_GRPO_synthetic_dataset.ipynb (GRPO + random Toy-IR corpus).""" +import json +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +OUT = ROOT / "Compilertetris_GRPO_synthetic_dataset.ipynb" + + +def cell_md(s: str) -> dict: + return { + "cell_type": "markdown", + "metadata": {}, + "source": [line + "\n" for line in s.strip().split("\n")], + } + + +def cell_code(s: str) -> dict: + return { + "cell_type": "code", + "metadata": {}, + "execution_count": None, + "outputs": [], + "source": [line + "\n" for line in s.rstrip().split("\n")], + } + +cells: list = [] + +cells.append(cell_md(""" +# Compiler Tetris β€” GRPO with **synthetic Toy-IR corpus** (Colab) + +This variant builds **50–200+** training prompts from `program_generator.py` (random const/add/mul chains in the same JSON shape as `runtime_core` / Deliverable 2). The original `metahack1` notebook uses a *different* IR schema; that generator is not mixed in here unless you add a converter. + +| What | Value | +|------|--------| +| Space | `greedybeserk95/Compilertetris` | +| Code root after clone | `/content/Compilertetris` | +| LoRA output | `/content/compilertetris_lora` | +| This run uses **shorter** `TRAIN_STEPS` than the 3-example notebook (see config cell). | +| **Checkpoints** | Under `output_dir` as `checkpoint-*` (see `CHECKPOINT_EVERY` / `KEEP_LAST_N_CHECKPOINTS` in the GRPO cell) | + +**Tuning:** In the dataset cell, set `N_TRAIN_PROGRAMS` (e.g. 80, 120, 200). +""")) + +cells.append( + cell_md( + """ +## Checkpoints + +The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `/content/grpo_compilertetris/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training). +""" + ) +) + +cells.append(cell_code(""" +# --- Central path config --- + +HF_SPACE_REPO = "https://huggingface.co/spaces/greedybeserk95/Compilertetris" +HF_ADAPTER_REPO = "greedybeserk95/Compilertetris-grpo-adapter" +REPO_DIR = "/content/Compilertetris" +""")) + +cells.append(cell_code(""" +import os, subprocess + +if os.path.isdir(REPO_DIR + "/.git"): + subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"], timeout=300) +else: + subprocess.check_call(["git", "clone", HF_SPACE_REPO, REPO_DIR], timeout=600) +print("Repo:", REPO_DIR) +""")) + +cells.append(cell_code("""!nvidia-smi""")) + +cells.append(cell_code(""" +!pip install -q "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" +!pip install -q trl datasets transformers accelerate peft bitsandbytes +""")) + +cells.append(cell_code(""" +import sys, json, os +sys.path.insert(0, REPO_DIR) + +import torch +import unsloth +from unsloth import FastLanguageModel +from datasets import Dataset + +from runtime_core import ( + CompilerOptimizationEnv, + Deliverable2_Formatter, + MockEngine, + MOCK_PASSES, +) +from program_generator import build_training_program_corpus + +print("torch", torch.__version__, "cuda", torch.cuda.is_available()) +print("MOCK_PASSES", list(MOCK_PASSES.keys())) +""")) + + +def _cell_chatml_synth() -> str: + p0 = """# Qwen2.5 ChatML + synthetic PROGRAMS +IM_END = \"<\" + \"|im_end|>\" +""" + p1 = r''' +import re + +# How many training programs (50–200 recommended; includes 3 builtins if include_builtins=True) +N_TRAIN_PROGRAMS = 120 +RANDOM_SEED = 42 + +SYSTEM_PROMPT = f"""You are a compiler optimization agent. +You will see Toy-IR as pseudo-assembly (Deliverable-2 text). + +Output ONLY a valid JSON array of optimization pass names IN ORDER. +Available passes: {", ".join(repr(p) for p in sorted(MOCK_PASSES))} +Rules: JSON array only; max 8 passes; you may repeat passes.""" + + +def build_prompt(program_list: list, program_id: int) -> str: + obs = Deliverable2_Formatter.translate_state(program_list) + return ( + f"<|im_start|>system\n{SYSTEM_PROMPT}{IM_END}\n" + f"<|im_start|>user\n#PROGRAM_ID:{program_id}\n{obs}{IM_END}\n" + f"<|im_start|>assistant\n" + ) + + +def program_from_prompt(prompt: str) -> list: + m = re.search(r"#PROGRAM_ID:(\d+)", prompt) + if m: + i = int(m.group(1)) + if 0 <= i < len(PROGRAMS): + return PROGRAMS[i] + if "#PROGRAM_JSON" in prompt: + tail = prompt.split("#PROGRAM_JSON", 1)[1] + if IM_END in tail: + tail = tail.split(IM_END, 1)[0] + raw = tail.strip() + if raw: + return json.loads(raw) + raise ValueError("cannot resolve program from prompt (expect #PROGRAM_ID:N or #PROGRAM_JSON)") + + +PROGRAMS = build_training_program_corpus( + n_total=N_TRAIN_PROGRAMS, + seed=RANDOM_SEED, + include_builtins=True, +) +print("Corpus size:", len(PROGRAMS)) + +train_dataset = Dataset.from_dict({ + "prompt": [build_prompt(p, i) for i, p in enumerate(PROGRAMS)], +}) +print("Dataset rows:", len(train_dataset)) +'''.lstrip("\n") + return p0 + p1 + + +cells.append(cell_code(_cell_chatml_synth())) + +cells.append(cell_code(""" +def env_reward_for_completion(prompt: str, completion: str, max_env_steps: int = 8) -> float: + try: + program = program_from_prompt(prompt) + except Exception: + return -8.0 + try: + actions = Deliverable2_Formatter.extract_action_array(completion) + except Exception: + return -5.0 + actions = [str(a).strip() for a in actions][: max_env_steps] + if not actions: + return -4.0 + env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=max_env_steps) + env.reset(program) + total = 0.0 + for a in actions: + step = env.step(a) + total += float(step.reward) + if step.done: + break + return float(total) + + +def make_reward_function(max_env_steps: int = 8): + def reward_func(prompts: list, completions: list, **kwargs) -> list: + return [env_reward_for_completion(p, c, max_env_steps) for p, c in zip(prompts, completions)] + return reward_func +""")) + +cells.append(cell_code(""" +from trl import GRPOConfig, GRPOTrainer + +MODEL_NAME = "unsloth/Qwen2.5-3B-Instruct" +MAX_SEQ_LEN = 1024 +MAX_COMPLETION = 256 + +model, tokenizer = FastLanguageModel.from_pretrained( + model_name=MODEL_NAME, + max_seq_length=MAX_SEQ_LEN, + dtype=None, + load_in_4bit=True, +) +model = FastLanguageModel.get_peft_model( + model, + r=16, lora_alpha=16, lora_dropout=0.0, bias="none", + use_gradient_checkpointing="unsloth", random_state=0, +) +""")) + +# Shorter run than 200 steps β€” scale up after smoke test +cells.append( + cell_code( + """ +NUM_GENERATIONS = 2 +LEARNING_RATE = 2e-5 +# Reduced from 200: increase after you verify loss/reward is stable +TRAIN_STEPS = 80 + +# --- Checkpoints (TrainingArguments / GRPO) --- +GRPO_OUTPUT_DIR = "/content/grpo_compilertetris" +CHECKPOINT_EVERY = 20 # save a checkpoint every N global steps +KEEP_LAST_N_CHECKPOINTS = 5 # on disk; older folders are deleted + +grpo_config = GRPOConfig( + output_dir=GRPO_OUTPUT_DIR, + learning_rate=LEARNING_RATE, + per_device_train_batch_size=NUM_GENERATIONS, + gradient_accumulation_steps=1, + num_generations=NUM_GENERATIONS, + max_completion_length=MAX_COMPLETION, + max_prompt_length=MAX_SEQ_LEN, + remove_unused_columns=False, + temperature=0.7, + max_steps=TRAIN_STEPS, + logging_steps=5, + save_strategy="steps", + save_steps=CHECKPOINT_EVERY, + save_total_limit=KEEP_LAST_N_CHECKPOINTS, + seed=0, + report_to="none", + use_vllm=False, +) + +reward_fn = make_reward_function() + +trainer = GRPOTrainer( + model=model, + processing_class=tokenizer, + reward_funcs=[reward_fn], + args=grpo_config, + train_dataset=train_dataset, +) +""" + ) +) + +cells.append(cell_code(""" +print("Starting GRPO (synthetic corpus)…") +trainer.train() +print("Done.") +""")) + +cells.append( + cell_code( + """ +# List on-disk checkpoints (for resume or manual export) +import glob, os + +ckpts = sorted( + glob.glob(os.path.join(GRPO_OUTPUT_DIR, "checkpoint-*")), + key=lambda p: int(p.split("checkpoint-")[-1]) if p.split("checkpoint-")[-1].isdigit() else 0, +) +print("Checkpoints in", GRPO_OUTPUT_DIR, ":", len(ckpts)) +for c in ckpts: + print(" ", c) +if ckpts: + print("Latest:", ckpts[-1]) +""" + ) +) + +cells.append( + cell_md( + """ +**Resume after disconnect / crash** β€” re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of: + +- `trainer.train(resume_from_checkpoint=True)` β€” continues from the latest `checkpoint-*` in `output_dir` +- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-40")` β€” specific step +""" + ) +) + +cells.append( + cell_code( + """ +# Uncomment to resume from the latest checkpoint (run after re-creating `trainer` in a new session) +# trainer.train(resume_from_checkpoint=True) +""" + ) +) + +cells.append( + cell_code( + """ +SAVE_DIR = "/content/compilertetris_lora" +trainer.model.save_pretrained(SAVE_DIR) +tokenizer.save_pretrained(SAVE_DIR) +print("Saved to", SAVE_DIR) +""" + ) +) + +cells.append( + cell_code( + """ +from huggingface_hub import login, HfApi + +login() +HfApi().create_repo(HF_ADAPTER_REPO, exist_ok=True, repo_type="model") +trainer.model.push_to_hub(HF_ADAPTER_REPO, private=True) +tokenizer.push_to_hub(HF_ADAPTER_REPO, private=True) +print("Pushed to https://huggingface.co/" + HF_ADAPTER_REPO) +""" + ) +) + +if __name__ == "__main__": + nb = { + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "colab": {"provenance": [], "gpuType": "T4"}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3", + }, + "language_info": {"name": "python"}, + }, + "cells": cells, + } + OUT.write_text(json.dumps(nb, indent=2), encoding="utf-8") + print("Wrote", OUT) diff --git a/space/space/space/space/write_colab_synth_notebook.py b/space/space/space/space/write_colab_synth_notebook.py index 5149477f96f750c894d3f896a70288afb9103b15..052f52201da8a5e76aa5622921cfcf22d42e1391 100644 --- a/space/space/space/space/write_colab_synth_notebook.py +++ b/space/space/space/space/write_colab_synth_notebook.py @@ -26,19 +26,21 @@ def cell_code(s: str) -> dict: cells: list = [] cells.append(cell_md(""" -# Compiler Tetris β€” GRPO with **synthetic Toy-IR corpus** (Colab) +# Compiler Tetris β€” GRPO, **synthetic Toy-IR** (Colab **T4 ~8–10 h** preset) -This variant builds **50–200+** training prompts from `program_generator.py` (random const/add/mul chains in the same JSON shape as `runtime_core` / Deliverable 2). The original `metahack1` notebook uses a *different* IR schema; that generator is not mixed in here unless you add a converter. +**Runtime (estimate):** after ~20 steps, read `it/s` in the progress bar. Approximate hours β‰ˆ `TRAIN_STEPS / (it/s * 3600)`. +With **`TRAIN_STEPS = 3600`**: at **0.10 it/s** that is about **10 h**; at **0.12 it/s** about **8.3 h**. If you finish much faster/slower, change **`TRAIN_STEPS`** (or the model size / `NUM_GENERATIONS`). + +**T4 tips:** 3B + 4bit + below settings fit T4; **7B** or very long `MAX_COMPLETION` can OOM. **High-RAM** runtime in Colab helps. Use **checkpoints** (next cells) in case the session dies. | What | Value | |------|--------| | Space | `greedybeserk95/Compilertetris` | -| Code root after clone | `/content/Compilertetris` | -| LoRA output | `/content/compilertetris_lora` | -| This run uses **shorter** `TRAIN_STEPS` than the 3-example notebook (see config cell). | -| **Checkpoints** | Under `output_dir` as `checkpoint-*` (see `CHECKPOINT_EVERY` / `KEEP_LAST_N_CHECKPOINTS` in the GRPO cell) | +| Code root | `/content/Compilertetris` | +| LoRA out | `/content/compilertetris_lora` | +| **Checkpoints** | `/content/grpo_compilertetris/checkpoint-*` | -**Tuning:** In the dataset cell, set `N_TRAIN_PROGRAMS` (e.g. 80, 120, 200). +`program_generator` β€” Toy-IR in `runtime_core` shape; `metahack1` uses a different schema. """)) cells.append( @@ -105,8 +107,8 @@ IM_END = \"<\" + \"|im_end|>\" p1 = r''' import re -# How many training programs (50–200 recommended; includes 3 builtins if include_builtins=True) -N_TRAIN_PROGRAMS = 120 +# T4: 200–300 is a good tradeoff (RAM + diversity) +N_TRAIN_PROGRAMS = 250 RANDOM_SEED = 42 SYSTEM_PROMPT = f"""You are a compiler optimization agent. @@ -183,7 +185,7 @@ def env_reward_for_completion(prompt: str, completion: str, max_env_steps: int = return float(total) -def make_reward_function(max_env_steps: int = 8): +def make_reward_function(max_env_steps: int = 10): def reward_func(prompts: list, completions: list, **kwargs) -> list: return [env_reward_for_completion(p, c, max_env_steps) for p, c in zip(prompts, completions)] return reward_func @@ -192,6 +194,7 @@ def make_reward_function(max_env_steps: int = 8): cells.append(cell_code(""" from trl import GRPOConfig, GRPOTrainer +# --- T4-friendly (3B 4-bit). For A100+ you can try Qwen2.5-7B and MAX_COMPLETION=384. --- MODEL_NAME = "unsloth/Qwen2.5-3B-Instruct" MAX_SEQ_LEN = 1024 MAX_COMPLETION = 256 @@ -204,37 +207,40 @@ model, tokenizer = FastLanguageModel.from_pretrained( ) model = FastLanguageModel.get_peft_model( model, - r=16, lora_alpha=16, lora_dropout=0.0, bias="none", + r=32, lora_alpha=32, lora_dropout=0.0, bias="none", use_gradient_checkpointing="unsloth", random_state=0, ) +print("Model:", MODEL_NAME, "| max_seq", MAX_SEQ_LEN, "| completion cap", MAX_COMPLETION) """)) -# Shorter run than 200 steps β€” scale up after smoke test +# ---- T4 / ~8–10 h wall time: tune TRAIN_STEPS after you see it/s in the first minutes ---- cells.append( cell_code( """ -NUM_GENERATIONS = 2 -LEARNING_RATE = 2e-5 -# Reduced from 200: increase after you verify loss/reward is stable -TRAIN_STEPS = 80 +# 4 rollouts per prompt: good for GRPO on T4; 6–8 is heavier (slower, more VRAM) +NUM_GENERATIONS = 4 +LEARNING_RATE = 1.5e-5 +# Target ~8–10 h on T4 when it/s is ~0.10–0.12 (typical for this stack). Re-tune if your it/s differs. +# hours β‰ˆ TRAIN_STEPS / (it/s * 3600) +TRAIN_STEPS = 3600 +GRAD_ACCUM = 2 -# --- Checkpoints (TrainingArguments / GRPO) --- GRPO_OUTPUT_DIR = "/content/grpo_compilertetris" -CHECKPOINT_EVERY = 20 # save a checkpoint every N global steps -KEEP_LAST_N_CHECKPOINTS = 5 # on disk; older folders are deleted +CHECKPOINT_EVERY = 200 +KEEP_LAST_N_CHECKPOINTS = 3 grpo_config = GRPOConfig( output_dir=GRPO_OUTPUT_DIR, learning_rate=LEARNING_RATE, per_device_train_batch_size=NUM_GENERATIONS, - gradient_accumulation_steps=1, + gradient_accumulation_steps=GRAD_ACCUM, num_generations=NUM_GENERATIONS, max_completion_length=MAX_COMPLETION, max_prompt_length=MAX_SEQ_LEN, remove_unused_columns=False, temperature=0.7, max_steps=TRAIN_STEPS, - logging_steps=5, + logging_steps=20, save_strategy="steps", save_steps=CHECKPOINT_EVERY, save_total_limit=KEEP_LAST_N_CHECKPOINTS, @@ -287,7 +293,7 @@ cells.append( **Resume after disconnect / crash** β€” re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of: - `trainer.train(resume_from_checkpoint=True)` β€” continues from the latest `checkpoint-*` in `output_dir` -- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-40")` β€” specific step +- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-300")` β€” example path (use an existing `checkpoint-*` folder) """ ) ) diff --git a/write_colab_synth_notebook.py b/write_colab_synth_notebook.py index 052f52201da8a5e76aa5622921cfcf22d42e1391..a65d9701b04a82eafd17fe4c559c271b06c8a0a6 100644 --- a/write_colab_synth_notebook.py +++ b/write_colab_synth_notebook.py @@ -36,11 +36,13 @@ With **`TRAIN_STEPS = 3600`**: at **0.10 it/s** that is about **10 h**; at **0.1 | What | Value | |------|--------| | Space | `greedybeserk95/Compilertetris` | -| Code root | `/content/Compilertetris` | -| LoRA out | `/content/compilertetris_lora` | -| **Checkpoints** | `/content/grpo_compilertetris/checkpoint-*` | +| Code root | `REPO_DIR` (printed in the β€œpaths” cell) | +| LoRA out | `LORA_DIR` (printed in the β€œpaths” cell) | +| **Checkpoints** | `GRPO_OUTPUT_DIR + "/checkpoint-*"` | `program_generator` β€” Toy-IR in `runtime_core` shape; `metahack1` uses a different schema. + +**Training evidence (loss + reward plots):** the cell *after* `trainer.train()` saves `training_loss_and_reward.png` under `output_dir` and shows it in the notebook β€” use it in your README / writeup. """)) cells.append( @@ -48,21 +50,49 @@ cells.append( """ ## Checkpoints -The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `/content/grpo_compilertetris/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training). +The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `$GRPO_OUTPUT_DIR/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training). """ ) ) cells.append(cell_code(""" -# --- Central path config --- +# --- Central path config (works on Colab *and* local / HF runtimes) --- + +import os +import tempfile +from pathlib import Path + +def _default_workspace_base() -> str: + # Colab: /content is writable + c = "/content" + if os.path.isdir(c) and os.access(c, os.W_OK): + return str(Path(c) / "work") + + # Otherwise: a guaranteed-writable temp dir + return str(Path(tempfile.gettempdir()) / "compilertetris_work") + + +BASE = os.environ.get("COMPILERTETRIS_BASE", _default_workspace_base()) +os.makedirs(BASE, exist_ok=True) HF_SPACE_REPO = "https://huggingface.co/spaces/greedybeserk95/Compilertetris" HF_ADAPTER_REPO = "greedybeserk95/Compilertetris-grpo-adapter" -REPO_DIR = "/content/Compilertetris" +REPO_DIR = str(Path(BASE) / "Compilertetris") +GRPO_OUTPUT_DIR = str(Path(BASE) / "grpo_compilertetris") +LORA_DIR = str(Path(BASE) / "compilertetris_lora") + +print("BASE :", BASE) +print("REPO_DIR :", REPO_DIR) +print("GRPO_OUTPUT :", GRPO_OUTPUT_DIR) +print("LORA_DIR :", LORA_DIR) """)) cells.append(cell_code(""" import os, subprocess +from pathlib import Path + +# Ensure parent is writable/created +Path(REPO_DIR).parent.mkdir(parents=True, exist_ok=True) if os.path.isdir(REPO_DIR + "/.git"): subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"], timeout=300) @@ -74,16 +104,22 @@ print("Repo:", REPO_DIR) cells.append(cell_code("""!nvidia-smi""")) cells.append(cell_code(""" -!pip install -q "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" -!pip install -q trl datasets transformers accelerate peft bitsandbytes +# One-shot install (A100 + CUDA 12.x): +# - Unsloth's `cu124-ampere-torch260` extra pulls a **Torch 2.6** CUDA stack (fixes `torch.int1` / torchao mismatches) +# - Pin TRL+Transformers to a Unsloth-tested pair (per Unsloth issue threads around torch2.6 / transformers4.55.x) +!python -m pip install -U --no-cache-dir \ + "transformers==4.55.4" "trl==0.20.0" \ + matplotlib datasets accelerate peft bitsandbytes safetensors huggingface_hub sentencepiece \ + "unsloth[cu124-ampere-torch260] @ git+https://github.com/unslothai/unsloth.git" """)) cells.append(cell_code(""" +# `import unsloth` should come before `transformers` is imported (happens via other imports). +import unsloth import sys, json, os sys.path.insert(0, REPO_DIR) import torch -import unsloth from unsloth import FastLanguageModel from datasets import Dataset @@ -225,7 +261,7 @@ LEARNING_RATE = 1.5e-5 TRAIN_STEPS = 3600 GRAD_ACCUM = 2 -GRPO_OUTPUT_DIR = "/content/grpo_compilertetris" +# GRPO_OUTPUT_DIR is set in the "Central path config" cell CHECKPOINT_EVERY = 200 KEEP_LAST_N_CHECKPOINTS = 3 @@ -268,6 +304,101 @@ trainer.train() print("Done.") """)) +cells.append( + cell_code( + r""" +# --- Loss + reward plots (submission): from trainer.state.log_history after train() --- + +!pip install -q matplotlib + +import os +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +PLOT_DIR = GRPO_OUTPUT_DIR +os.makedirs(PLOT_DIR, exist_ok=True) +PNG_PATH = os.path.join(PLOT_DIR, "training_loss_and_reward.png") + +def _reward_from_log_row(h): + v = h.get("reward") + if isinstance(v, (int, float)): + return float(v), "reward" + best_k, best_v = None, None + for k, v in h.items(): + if not isinstance(v, (int, float)): + continue + klow = k.replace("-", "/").lower() + if "reward" not in klow or "std" in klow or "clip" in klow or "per_token" in klow: + continue + if "mean" in klow or k == "reward": + return float(v), k + if "mean" not in klow and best_k is None: + best_k, best_v = k, float(v) + if best_k is not None: + return best_v, best_k + return None, None + +def extract_series(history): + sl, vl, sr, vr = [], [], [], [] + rlabel = None + for h in history: + s = h.get("step") + if s is None: + continue + lo = h.get("loss") + if isinstance(lo, (int, float)): + sl.append(s) + vl.append(float(lo)) + r_val, rk = _reward_from_log_row(h) + if r_val is not None and rk: + if rlabel is None: + rlabel = rk + if rk == rlabel: + sr.append(s) + vr.append(r_val) + return (sl, vl, "loss"), (sr, vr, rlabel or "reward") + +(loss_s, loss_v, _lk), (rew_s, rew_v, rew_lab) = extract_series(trainer.state.log_history) +print("Points β€” loss:", len(loss_v), "| reward:", len(rew_v)) +if trainer.state.log_history: + print("Last log row keys (sample):", list(trainer.state.log_history[-1].keys())[:25]) + +fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(10, 6), sharex=True) +if loss_v: + ax0.plot(loss_s, loss_v, "b.-", label="loss", linewidth=1, markersize=2) + ax0.set_ylabel("training loss") + ax0.set_title("GRPO (this Colab run)") + ax0.grid(True, alpha=0.3) + ax0.legend() +else: + ax0.text(0.5, 0.5, "No 'loss' in log_history", ha="center", transform=ax0.transAxes) + +if rew_v: + ax1.plot(rew_s, rew_v, "g.-", label=rew_lab, linewidth=1, markersize=2) + ax1.set_ylabel("mean reward" if "mean" in (rew_lab or "") else "reward") + ax1.set_xlabel("global step") + ax1.grid(True, alpha=0.3) + ax1.legend() +else: + ax1.text(0.5, 0.5, "No reward column found β€” see keys above", ha="center", transform=ax1.transAxes) + ax1.set_xlabel("global step") + +plt.tight_layout() +plt.savefig(PNG_PATH, dpi=150, bbox_inches="tight") +print("Saved:", PNG_PATH) +try: + from IPython.display import Image, display + display(Image(PNG_PATH)) +except Exception as e: + print("Display:", e) +finally: + plt.close("all") +print("If keys differ: print(trainer.state.log_history[-1])") +""" + ) +) + cells.append( cell_code( """ @@ -293,7 +424,7 @@ cells.append( **Resume after disconnect / crash** β€” re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of: - `trainer.train(resume_from_checkpoint=True)` β€” continues from the latest `checkpoint-*` in `output_dir` -- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-300")` β€” example path (use an existing `checkpoint-*` folder) +- `trainer.train(resume_from_checkpoint="/checkpoint-300")` β€” example (use an existing `checkpoint-*` folder; print `GRPO_OUTPUT_DIR` from the path cell) """ ) ) @@ -310,7 +441,7 @@ cells.append( cells.append( cell_code( """ -SAVE_DIR = "/content/compilertetris_lora" +SAVE_DIR = LORA_DIR trainer.model.save_pretrained(SAVE_DIR) tokenizer.save_pretrained(SAVE_DIR) print("Saved to", SAVE_DIR)