refactor: Builder System

examples/win_at_p.py

@@ -26,7 +26,6 @@
 Notes:
 - If multiple runs exist for the same author within a group, we take the MIN latency for that author in that group.
 - By default, the baseline author ('flashinfer') is EXCLUDED from output curves; use --include-baseline to include it.
-
 """
 
 import argparse

flashinfer_bench/apply/key.py

@@ -12,7 +12,7 @@
 class ApplyKey:
     axes: Tuple[Tuple[str, int], ...] = field(default_factory=tuple)
     # Features extracted from input tensors
-    feats: Tuple[Tuple[str, Union[int, Union[float, bool]]], ...] = field(default_factory=tuple)
+    feats: Tuple[Tuple[str, Union[int, float, bool]], ...] = field(default_factory=tuple)
 
     def encode(self) -> str:
         return json.dumps(

flashinfer_bench/compile/__init__.py

@@ -1,10 +1,27 @@
 """Compiler subsystem package.
 
-Exports common builder types for convenience.
+This package provides the infrastructure for building solutions into executable runnables.
+It includes:
+- Builder: Abstract base class for different language/build system implementations
+- BuilderRegistry: Central registry for managing and dispatching builders
+- Runnable: Executable wrapper around compiled solutions
+- RunnableMetadata: Metadata about build process and source
+
+The typical workflow is:
+1. Get the singleton registry: registry = BuilderRegistry.get_instance()
+2. Build a solution: runnable = registry.build(definition, solution)
+3. Execute: result = runnable(**inputs)
 """
 
 from .builder import Builder, BuildError
 from .registry import BuilderRegistry, get_builder_registry
-from .runnable import Runnable
+from .runnable import Runnable, RunnableMetadata
 
-__all__ = ["Builder", "BuildError", "BuilderRegistry", "Runnable", "get_builder_registry"]
+__all__ = [
+    "Builder",
+    "BuildError",
+    "BuilderRegistry",
+    "Runnable",
+    "RunnableMetadata",
+    "get_builder_registry",
+]

flashinfer_bench/compile/builder.py

@@ -1,95 +1,93 @@
+"""Abstract base class for solution builders."""
+
 from __future__ import annotations
 
-import hashlib
-import os
-import re
-import tempfile
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Optional
-
-from flashinfer_bench.compile.runnable import Runnable
-from flashinfer_bench.data import Definition, Solution, SourceFile
-
-
-def write_sources_to_dir(dir: str, sources: list[SourceFile]) -> None:
-    os.makedirs(dir, exist_ok=True)
-    for src in sources:
-        abspath = os.path.join(dir, src.path)
-        os.makedirs(os.path.dirname(abspath), exist_ok=True)
-        with open(abspath, "w", encoding="utf-8") as f:
-            f.write(src.content)
-
-
-def write_sources_to_temp(base: str, sources: list[SourceFile], pkg: Optional[str] = None) -> str:
-    os.makedirs(base, exist_ok=True)
-    tmpdir = tempfile.mkdtemp(dir=base)
-    if pkg:
-        tmpdir = os.path.join(tmpdir, pkg)
-        os.makedirs(tmpdir, exist_ok=True)
-    write_sources_to_dir(tmpdir, sources)
-    return tmpdir
-
-
-def create_pkg_name(sol: Solution, prefix: str = "") -> str:
-    # Normalize the solution name
-    s = re.sub(r"[^0-9a-zA-Z_]", "_", sol.name)
-    if not s or s[0].isdigit():
-        s = "_" + s
+from pathlib import Path
+from typing import Dict, Tuple
 
-    # Hash the sources
-    h = hashlib.sha1()
-    for src in sol.sources:
-        h.update(src.path.encode())
-        h.update(src.content.encode())
+from flashinfer_bench.data import Definition, Solution
+from flashinfer_bench.env import get_fib_cache_path
 
-    return prefix + s + "_" + h.hexdigest()[:6]
+from .runnable import Runnable
+from .utils import create_package_name
 
 
 class BuildError(RuntimeError):
     """Raised when a builder fails to construct a runnable implementation."""
 
 
 class Builder(ABC):
-    """Builder abstraction: (Definition, Solution) -> Runnable with hidden cache."""
+    """Abstract base class for building solutions into runnable implementations.
 
-    def __init__(self) -> None:
-        self._cache: Dict[str, Runnable] = {}
+    A Builder transforms a (Definition, Solution) pair into a Runnable object, which is an
+    executable implementation of the solution. Different builders handle different programming
+    languages (e.g., Python, CUDA, Triton) and build systems.
 
-    @abstractmethod
-    def can_build(self, solution: Solution) -> bool:
-        """Build guard to check if this builder can handle the given solution."""
-        ...
+    Subclasses must implement all its abstract methods. Expectedly, the concrete builder should
+    operate in the folder `FIB_CACHE_PATH / builder_specific_subfolder / key`, where `key` is
+    a unique identifier for the solution.
+    """
 
-    @abstractmethod
-    def _build(self, definition: Definition, solution: Solution) -> Runnable:
-        """Perform a real build and return a Runnable; raise BuildError on failure."""
-        ...
+    def __init__(self, key_prefix: str, build_dir_name: str) -> None:
+        """Initialize the builder."""
+        self._key_prefix = key_prefix
+        self._build_dir_name = build_dir_name
 
     @abstractmethod
-    def _make_closer(self, *args, **kwargs) -> Callable[[], None]:
-        """Factory for a resource closer used by the concrete builder."""
+    def can_build(self, solution: Solution) -> bool:
+        """Check if this builder can handle the given solution.
+
+        Parameters
+        ----------
+        solution : Solution
+            The solution to check.
+
+        Returns
+        -------
+        bool
+            True if this builder can build the solution, False otherwise.
+        """
         ...
 
     @abstractmethod
-    def _make_key(self, solution: Solution) -> str:
-        """Cache key for a solution."""
+    def build(self, definition: Definition, solution: Solution) -> Runnable:
+        """Build a solution into a runnable implementation.
+
+        This method compiles/loads the solution's source code and returns a Runnable
+        object that can be executed with the interface specified by the definition.
+
+        Parameters
+        ----------
+        definition : Definition
+            The problem definition that specifies the expected interface.
+        solution : Solution
+            The solution implementation to build.
+
+        Returns
+        -------
+        Runnable
+            An executable wrapper around the built implementation.
+
+        Raises
+        ------
+        BuildError
+            If the build fails for any reason (compilation errors, missing dependencies, etc.).
+        """
         ...
 
-    def build(self, definition: Definition, solution: Solution) -> Runnable:
-        """Public entry with per-solution cache keyed by solution.name."""
-        key = self._make_key(solution)
-        if key in self._cache:
-            return self._cache[key]
-        runnable = self._build(definition, solution)
-        self._cache[key] = runnable
-        return runnable
-
-    def clear_cache(self) -> None:
-        """Close all cached runnables and clear the cache."""
-        for r in list(self._cache.values()):
-            try:
-                r.close()
-            except Exception:
-                # Best-effort cleanup; keep going
-                pass
-        self._cache.clear()
+    def get_package_name_and_build_path(self, solution: Solution) -> Tuple[str, Path]:
+        """Get the package name and build path for the solution.
+
+        Parameters
+        ----------
+        solution : Solution
+            The solution to get the package name and build path for.
+
+        Returns
+        -------
+        Tuple[str, Path]: The package name and build path for the solution.
+        """
+        package_name = create_package_name(solution, self._key_prefix)
+        build_path = get_fib_cache_path() / self._build_dir_name / package_name
+        return package_name, build_path

flashinfer_bench/compile/builders/__init__.py

@@ -1,6 +1,8 @@
-from .cuda_builder import CUDABuilder
+"""Concrete builder implementations for different languages and build systems."""
+
 from .python_builder import PythonBuilder
+from .torch_builder import TorchBuilder
 from .triton_builder import TritonBuilder
 from .tvm_ffi_builder import TVMFFIBuilder
 
-__all__ = ["CUDABuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]
+__all__ = ["TorchBuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"]