Class KernelBenchProblem

kernelbench/bridge.py:48–347

A single KernelBench problem.

@dataclass
class KernelBenchProblem:
    """A single KernelBench problem."""

    level: int
    problem_id: int
    name: str
    source_code: str

    @property
    def uid(self) -> str:
        return f"L{self.level}_P{self.problem_id:03d}"

    @property
    def cache_path(self) -> Path:
        return KB_CACHE_DIR / f"level{self.level}" / f"{self.problem_id}.py"

    # ----- Cache persistence -----

    def save_to_cache(self) -> None:
        self.cache_path.parent.mkdir(parents=True, exist_ok=True)
        self.cache_path.write_text(self.source_code, encoding="utf-8")
        meta_path = self.cache_path.with_suffix(".json")
        meta = {
            "level": self.level,
            "problem_id": self.problem_id,
            "name": self.name,
        }
        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")

    @classmethod
    def load_from_cache(cls, level: int, problem_id: int) -> Optional["KernelBenchProblem"]:
        cache_path = KB_CACHE_DIR / f"level{level}" / f"{problem_id}.py"
        meta_path = cache_path.with_suffix(".json")
        if not cache_path.exists():
            return None
        source = cache_path.read_text(encoding="utf-8")
        name = f"problem_{problem_id}"
        if meta_path.exists():
            meta = json.loads(meta_path.read_text(encoding="utf-8"))
            name = meta.get("name", name)
        return cls(level=level, problem_id=problem_id, name=name, source_code=source)

    # ----- Analysis -----

    def analyze(self) -> Dict[str, Any]:
        """Identify operations, shapes, parameter usage, and estimate difficulty."""
        analysis: Dict[str, Any] = {
            "operations": [],
            "estimated_difficulty": "unknown",
            "has_parameters": False,
            "input_shapes": [],
            "forward_lines": 0,
        }

        try:
            ast.parse(self.source_code)
        except SyntaxError:
            return analysis