Function setup_problem

kernelbench/bridge.py:495–547

Set up workspace for optimizing a KernelBench problem. Creates: workspace/kb_active/reference.py -- original Model + get_inputs workspace/kb_active/metadata.json -- problem metadata + analysis kernel.py -- starter ModelNew (edit this)

(problem: KernelBenchProblem, backend: str = "cuda")

# ---------------------------------------------------------------------------

def setup_problem(problem: KernelBenchProblem, backend: str = "cuda") -> None:
    """
    Set up workspace for optimizing a KernelBench problem.

    Creates:
      workspace/kb_active/reference.py   -- original Model + get_inputs
      workspace/kb_active/metadata.json  -- problem metadata + analysis
      kernel.py                          -- starter ModelNew (edit this)
    """
    KB_ACTIVE_DIR.mkdir(parents=True, exist_ok=True)

    # Write reference
    ref_path = KB_ACTIVE_DIR / "reference.py"
    ref_path.write_text(problem.source_code, encoding="utf-8")

    # Write metadata
    meta_path = KB_ACTIVE_DIR / "metadata.json"
    analysis = problem.analyze()
    metadata = {
        "level": problem.level,
        "problem_id": problem.problem_id,
        "name": problem.name,
        "uid": problem.uid,
        "analysis": analysis,
        "backend": backend,
    }
    meta_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

    # Generate starter kernel.py
    starter = problem.generate_starter(backend=backend)
    KERNEL_PY.write_text(starter, encoding="utf-8")

    # Report
    print(f"=== KernelBench Problem Setup ===")
    print()
    print(f"  Problem:    {problem.uid} -- {problem.name}")
    print(f"  Level:      {problem.level}")
    print(f"  Operations: {', '.join(analysis['operations']) or 'unknown'}")
    print(f"  Difficulty: {analysis['estimated_difficulty']}")
    print(f"  Parameters: {'yes' if analysis['has_parameters'] else 'no'}")
    print(f"  Forward:    {analysis['forward_lines']} lines")
    if analysis["input_shapes"]:
        for i, s in enumerate(analysis["input_shapes"]):
            print(f"  Input {i}:    shape={s}")
    print()
    print(f"  Reference:  workspace/kb_active/reference.py")
    print(f"  Metadata:   workspace/kb_active/metadata.json")
    print(f"  Kernel:     kernel.py  <-- EDIT THIS")
    print()
    print("Next steps:")
    print("  1. Edit kernel.py -- optimize ModelNew.forward()")
    print("  2. Run: uv run kernelbench/bench_kb.py")
    print("  3. Keep improvements or git reset --hard HEAD~1 to revert")


# ---------------------------------------------------------------------------