Use LLM/KernelAgent backends for any op suite (#63)

Author: bertmaher

BackendBench/scripts/main.py

@@ -84,6 +84,12 @@ def setup_logging(log_level):
     type=str,
     help="Path to TorchBench operator data",
 )
+@click.option(
+    "--ops-directory",
+    default="generated_kernels",
+    type=str,
+    help="Path to directory containing generated kernels",
+)
 def cli(
     log_level,
     suite,
@@ -94,6 +100,7 @@ def cli(
     kernel_agent_workers,
     kernel_agent_max_rounds,
     torchbench_data_path,
+    ops_directory,
 ):
     setup_logging(log_level)
     if ops:
@@ -107,17 +114,6 @@ def cli(
         "directory": backends.DirectoryBackend,
     }[backend]()
 
-    # For LLM backend, we need to generate kernels first
-    if backend.name == "llm":
-        llm_client = ClaudeKernelGenerator()
-        backend = setup_llm_backend(backend, llm_client, suite, ops, llm_max_attempts)
-
-    # For KernelAgent backend, we need to generate kernels using the sophisticated agent system
-    elif backend.name == "kernel_agent":
-        backend = setup_kernel_agent_backend(
-            backend, suite, ops, kernel_agent_workers, kernel_agent_max_rounds
-        )
-
     suite = {
         "smoke": lambda: SmokeTestSuite,
         "opinfo": lambda: OpInfoTestSuite(
@@ -134,6 +130,21 @@ def cli(
         ),
     }[suite]()
 
+    # For LLM backend, we need to generate kernels first
+    if backend.name == "llm":
+        llm_client = ClaudeKernelGenerator()
+        backend = setup_llm_backend(backend, llm_client, suite, llm_max_attempts)
+
+    # For KernelAgent backend, we need to generate kernels using the sophisticated agent system
+    elif backend.name == "kernel_agent":
+        backend = setup_kernel_agent_backend(
+            backend, suite, kernel_agent_workers, kernel_agent_max_rounds
+        )
+
+    # For Directory backend, we need to load existing kernels from a directory
+    elif backend.name == "directory":
+        backend = backends.DirectoryBackend(ops_directory)
+
     overall_correctness = []
     overall_performance = []
 
@@ -160,21 +171,9 @@ def cli(
     print(f"performance score (geomean speedup over all operators): {geomean_perf:.2f}")
 
 
-def setup_llm_backend(llm_backend, llm_client, suite_name, ops_filter, max_attempts=5):
+def setup_llm_backend(llm_backend, llm_client, suite, max_attempts=5):
     """Setup LLM backend by generating kernels for all operations in the suite."""
     try:
-        if suite_name == "smoke":
-            suite = SmokeTestSuite
-        elif suite_name == "opinfo":
-            suite = OpInfoTestSuite(
-                "opinfo_cuda_bfloat16",
-                "cuda",
-                torch.bfloat16,
-                filter=ops_filter,
-            )
-        else:
-            raise ValueError(f"Unknown suite: {suite_name}")
-
         successful_ops = 0
         total_ops = 0
 
@@ -287,24 +286,11 @@ def feedback_callback(kernel_code: str, attempt: int) -> tuple[bool, Dict]:
         sys.exit(1)
 
 
-def setup_kernel_agent_backend(
-    kernel_agent_backend, suite_name, ops_filter, num_workers=4, max_rounds=10
-):
+def setup_kernel_agent_backend(kernel_agent_backend, suite, num_workers=4, max_rounds=10):
     """Setup KernelAgent backend by generating kernels using the sophisticated agent system."""
     try:
         # Configure the backend with the specified parameters
         kernel_agent_backend.set_config(num_workers, max_rounds)
-        if suite_name == "smoke":
-            suite = SmokeTestSuite
-        elif suite_name == "opinfo":
-            suite = OpInfoTestSuite(
-                "opinfo_cuda_bfloat16",
-                "cuda",
-                torch.bfloat16,
-                filter=ops_filter,
-            )
-        else:
-            raise ValueError(f"Unknown suite: {suite_name}")
 
         successful_ops = 0
         total_ops = 0