Directory Backend (#30)

Author: msaroufim

.gitignore

@@ -3,5 +3,6 @@ __pycache__/
 .vscode/
 .ruff_cache/
 generated_kernels/
+CLAUDE.md
 venv/
-CLAUDE.md
+ops/

BackendBench/backends.py

@@ -1,13 +1,98 @@
 import os
 import importlib.util
+import logging
 from typing import Dict, Callable, List
 
+logger = logging.getLogger(__name__)
+
 
 class Backend:
     def __init__(self, name):
         self.name = name
 
 
+class DirectoryBackend(Backend):
+    def __init__(self, ops_dir="generated_kernels"):
+        super().__init__("directory")
+        self.ops_dir = ops_dir
+        self.compiled_kernels: Dict[str, Callable] = {}
+        self._load_kernels()
+
+    def _load_kernels(self):
+        if not os.path.exists(self.ops_dir):
+            logger.warning(f"ops directory {self.ops_dir} does not exist")
+            return
+
+        loaded_count = 0
+        for op_name in os.listdir(self.ops_dir):
+            op_dir = os.path.join(self.ops_dir, op_name)
+            if not os.path.isdir(op_dir):
+                continue
+
+            impl_files = [f for f in os.listdir(op_dir) if f.endswith(".py")]
+            if not impl_files:
+                logger.warning(f"No Python files found in {op_dir}")
+                continue
+
+            # Use the first implementation file
+            impl_file = impl_files[0]
+            impl_path = os.path.join(op_dir, impl_file)
+
+            try:
+                # Load the implementation and map to PyTorch operation
+                kernel_func = self._load_kernel_from_file(impl_path, op_name)
+                pytorch_op = self._find_pytorch_op(op_name)
+                if pytorch_op:
+                    self.compiled_kernels[pytorch_op] = kernel_func
+                    logger.info(f"Loaded {op_name} from {impl_file}")
+                    loaded_count += 1
+                else:
+                    logger.warning(f"Could not map {op_name} to PyTorch operation")
+
+            except Exception as e:
+                logger.error(f"Error loading {op_name} from {impl_file}: {e}")
+
+        logger.info(f"DirectoryBackend loaded {loaded_count} kernels from {self.ops_dir}/")
+
+    def _load_kernel_from_file(self, file_path: str, op_name: str) -> Callable:
+        spec = importlib.util.spec_from_file_location(f"op_{op_name}", file_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        kernel_func_name = f"{op_name}_kernel_impl"
+        if hasattr(module, kernel_func_name):
+            return getattr(module, kernel_func_name)
+        else:
+            raise ValueError(f"No callable function found in {file_path}")
+
+    def _find_pytorch_op(self, op_name: str):
+        """Map operation name to PyTorch operation."""
+        import torch
+
+        # Try common patterns
+        try:
+            return getattr(torch.ops.aten, op_name).default
+        except AttributeError:
+            pass
+
+        try:
+            return getattr(torch.ops.aten, op_name).Tensor
+        except AttributeError:
+            pass
+
+        # Not 100% sure this is right, will need to iterate over all ops
+        return None
+
+    def __getitem__(self, key):
+        if key in self.compiled_kernels:
+            return self.compiled_kernels[key]
+        # Fallback to original operation if not implemented
+        return key
+
+    def __contains__(self, key):
+        return key in self.compiled_kernels or True  # Always claim to contain ops for fallback
+
+
 class AtenBackend(Backend):
     def __init__(self) -> None:
         super().__init__("aten")

README.md

@@ -73,3 +73,97 @@ Run KernelAgent on opinfo tests with a specific operation:
 export OPENAI_API_KEY=your_api_key_here
 python scripts/main.py --suite opinfo --backend kernel_agent --ops "add"
 ```
+
+## Directory-Based Kernel Development
+
+BackendBench supports a simple directory structure for manually adding kernel implementations. This is perfect for researchers who want to contribute optimized kernels without dealing with complex generation systems.
+
+### Directory Structure
+
+Create kernels in the following structure:
+```
+generated_kernels/
+├── relu/
+│   └── relu_implementation_1.py
+├── add/  
+│   └── add_implementation_1.py
+├── mul/
+│   └── mul_implementation_1.py
+└── ...
+```
+
+### How to Add Your Kernels
+
+1. **Create the operation directory:**
+   ```bash
+   mkdir generated_kernels/{op_name}
+   ```
+
+2. **Create your implementation file:**
+   ```bash
+   # Example: generated_kernels/relu/relu_implementation_1.py
+   ```
+
+3. **Write your kernel following this template:**
+   ```python
+   import torch
+   
+   def {op_name}_kernel_impl(*args, **kwargs):
+       """
+       Your kernel implementation.
+       Must match the PyTorch operation signature exactly.
+       """
+       # Your implementation here
+       return result
+   
+   # Optional: Add a test
+   if __name__ == "__main__":
+       pass
+   ```
+
+### Operation Name Mapping
+
+Use these exact directory names for common operations:
+- `relu` → `torch.ops.aten.relu.default`  
+- `add` → `torch.ops.aten.add.Tensor`
+- `mul` → `torch.ops.aten.mul.Tensor` 
+- `div` → `torch.ops.aten.div.Tensor`
+
+To find the correct name for other operations:
+```python
+# Find operation name
+import torch
+op = torch.ops.aten.some_op.some_variant
+print(str(op).split('aten.')[-1].split('.')[0])  # Use this as directory name
+```
+
+### Example Implementation
+
+Here's a complete example for ReLU:
+
+```python
+# generated_kernels/relu/relu_implementation_1.py
+import torch
+
+def relu_kernel_impl(input_tensor):
+    return torch.maximum(input_tensor, torch.zeros_like(input_tensor))
+
+if __name__ == "__main__":
+    # Test on CPU
+    x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
+    result = relu_kernel_impl(x)
+    expected = torch.tensor([0.0, 0.0, 0.0, 1.0, 2.0])
+    print(f"Test passed: {torch.allclose(result, expected)}")
+```
+
+### Testing Your Kernels
+
+Test individual implementations:
+```bash
+python generated_kernels/relu/relu_implementation_1.py
+```
+
+Test with BackendBench:
+```bash
+python scripts/main.py --suite smoke --backend directory
+```

requirements.txt

@@ -3,3 +3,4 @@ click
 numpy
 expecttest
 anthropic>=0.34.0
+pytest

scripts/create_simple_test_ops.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""
+Create simple kernel implementations for 5 common operations.
+Each just calls the original PyTorch function.
+"""
+
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def create_relu():
+    os.makedirs("generated_kernels/relu", exist_ok=True)
+    with open("generated_kernels/relu/relu_implementation_1.py", "w") as f:
+        f.write('''import torch
+
+def relu_kernel_impl(input):
+    """Simple ReLU implementation."""
+    return torch.ops.aten.relu.default(input)
+
+if __name__ == "__main__":
+    x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
+    result = relu_kernel_impl(x)
+    expected = torch.tensor([0.0, 0.0, 0.0, 1.0, 2.0])
+    print(f"ReLU test passed: {torch.allclose(result, expected)}")
+''')
+    logger.info("Created relu implementation")
+
+
+def create_add():
+    os.makedirs("generated_kernels/add", exist_ok=True)
+    with open("generated_kernels/add/add_implementation_1.py", "w") as f:
+        f.write('''import torch
+
+def add_kernel_impl(input, other):
+    """Simple add implementation."""
+    return torch.ops.aten.add.Tensor(input, other)
+
+if __name__ == "__main__":
+    a = torch.tensor([1.0, 2.0, 3.0])
+    b = torch.tensor([4.0, 5.0, 6.0])
+    result = add_kernel_impl(a, b)
+    expected = torch.tensor([5.0, 7.0, 9.0])
+    print(f"Add test passed: {torch.allclose(result, expected)}")
+''')
+    logger.info("Created add implementation")
+
+
+def create_mul():
+    os.makedirs("generated_kernels/mul", exist_ok=True)
+    with open("generated_kernels/mul/mul_implementation_1.py", "w") as f:
+        f.write('''import torch
+
+def mul_kernel_impl(input, other):
+    """Simple mul implementation."""
+    return torch.ops.aten.mul.Tensor(input, other)
+
+if __name__ == "__main__":
+    a = torch.tensor([1.0, 2.0, 3.0])
+    b = torch.tensor([4.0, 5.0, 6.0])
+    result = mul_kernel_impl(a, b)
+    expected = torch.tensor([4.0, 10.0, 18.0])
+    print(f"Mul test passed: {torch.allclose(result, expected)}")
+''')
+    logger.info("Created mul implementation")
+
+
+def create_abs():
+    os.makedirs("generated_kernels/abs", exist_ok=True)
+    with open("generated_kernels/abs/abs_implementation_1.py", "w") as f:
+        f.write('''import torch
+
+def abs_kernel_impl(input):
+    """Simple abs implementation."""
+    return torch.ops.aten.abs.default(input)
+
+if __name__ == "__main__":
+    x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
+    result = abs_kernel_impl(x)
+    expected = torch.tensor([2.0, 1.0, 0.0, 1.0, 2.0])
+    print(f"Abs test passed: {torch.allclose(result, expected)}")
+''')
+    logger.info("Created abs implementation")
+
+
+def create_sum():
+    os.makedirs("generated_kernels/sum", exist_ok=True)
+    with open("generated_kernels/sum/sum_implementation_1.py", "w") as f:
+        f.write('''import torch
+
+def sum_kernel_impl(input, *args, **kwargs):
+    """Simple sum implementation."""
+    return torch.ops.aten.sum.default(input, *args, **kwargs)
+
+if __name__ == "__main__":
+    x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    result = sum_kernel_impl(x)
+    expected = torch.tensor(10.0)
+    print(f"Sum test passed: {torch.allclose(result, expected)}")
+''')
+    logger.info("Created sum implementation")
+
+
+def main():
+    """Create 5 simple test operations."""
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    logger.info("Creating simple test implementations...")
+
+    create_relu()
+    create_add()
+    create_mul()
+    create_abs()
+    create_sum()
+
+    logger.info("Created 5 simple kernel implementations in generated_kernels/")
+    logger.info("Test them individually:")
+    logger.info("  python generated_kernels/relu/relu_implementation_1.py")
+    logger.info("  python generated_kernels/add/add_implementation_1.py")
+    logger.info("  etc.")
+    logger.info("Or test all with the backend:")
+    logger.info("  python test/test_simple_directory_backend.py")
+
+
+if __name__ == "__main__":
+    main()

scripts/main.py

@@ -44,7 +44,7 @@ def setup_logging(log_level):
 @click.option(
     "--backend",
     default="aten",
-    type=click.Choice(["aten", "flag_gems", "llm", "kernel_agent"]),
+    type=click.Choice(["aten", "flag_gems", "llm", "kernel_agent", "directory"]),
     help="Which backend to run",
 )
 @click.option(
@@ -96,6 +96,7 @@ def cli(
         "flag_gems": backends.FlagGemsBackend,
         "llm": backends.LLMBackend,
         "kernel_agent": backends.KernelAgentBackend,
+        "directory": backends.DirectoryBackend,
     }[backend]()
 
     # For LLM backend, we need to generate kernels first