Filter out solutions that have cuda streams (#56)

Author: PaliC

BackendBench/eval.py

@@ -5,6 +5,8 @@
 import triton.testing
 
 
+from BackendBench.utils import uses_cuda_stream
+
 logger = logging.getLogger(__name__)
 
 EXC_MSG = """
@@ -101,6 +103,11 @@ def eval_performance(op, impl, tests):
 
 def eval_one_op(op, impl, correctness_tests, performance_tests):
     """Evaluate impl of op against correctness_tests and performance_tests."""
+    # TODO: We should have proper error reporting instead of just saying this is 0,
+    # but that should be a separate PR.
+    if uses_cuda_stream(impl):
+        logger.warning(f"Skipping {op.__name__} because it uses CUDA stream")
+        return 0, 0
     return eval_correctness(op, impl, correctness_tests), eval_performance(
         op, impl, performance_tests
     )

BackendBench/utils.py

@@ -0,0 +1,53 @@
+import ast
+import inspect
+import re
+import textwrap
+
+
+def uses_cuda_stream(func) -> bool:
+    """
+    Detects whether a Python function creates CUDA streams.
+
+    Args:
+        func: The Python function to analyze
+
+    Returns:
+        bool: True if CUDA streams are created, False otherwise
+    """
+    try:
+        source = inspect.getsource(func)
+    except (TypeError, OSError):
+        # Handle builtin functions, OpOverload objects, and other callables
+        # without source code. These cannot create CUDA streams.
+        return False
+
+    # Check for stream creation patterns
+    patterns = [
+        r"torch\.cuda\.Stream\(",  # torch.cuda.Stream() constructor
+        r"cupy\.cuda\.Stream\(",  # cupy.cuda.Stream() constructor
+        r"cuda\.Stream\(",  # Generic cuda.Stream() constructor
+        r"pycuda.*Stream\(",  # PyCUDA stream creation
+        r"\bStream\(",  # Stream() constructor calls
+        r"make_stream\(",  # make_stream() factory function
+        r"create_stream\(",  # create_stream() factory function
+    ]
+
+    if any(re.search(p, source, re.IGNORECASE) for p in patterns):
+        return True
+
+    class StreamCreationFinder(ast.NodeVisitor):
+        def __init__(self):
+            self.found = False
+
+        def visit_Call(self, node):
+            # Check for Stream() constructor calls
+            if hasattr(node.func, "attr") and node.func.attr == "Stream":
+                self.found = True
+            elif hasattr(node.func, "id") and node.func.id == "Stream":
+                self.found = True
+            self.generic_visit(node)
+
+    tree = ast.parse(textwrap.dedent(source))
+    finder = StreamCreationFinder()
+    finder.visit(tree)
+    return finder.found

pyproject.toml

@@ -33,6 +33,9 @@ dev = [
     "pytest-mock",
     "pytest-timeout",
     "ruff==0.12.1",
+    "torch",
+    "numpy",
+    "cupy-cuda12x",
 ]
 flaggems = [
     "flag_gems",

test/test_utils.py

@@ -0,0 +1,163 @@
+import pytest
+from BackendBench.utils import uses_cuda_stream
+
+# Check if CUDA is available
+import torch
+
+HAS_CUDA = torch.cuda.is_available()
+
+
+class TestCudaStreamDetection:
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_pytorch_stream_creation(self):
+        """Test detection of PyTorch CUDA stream creation."""
+
+        def func_with_pytorch_stream():
+            import torch
+
+            stream = torch.cuda.Stream()
+            return stream
+
+        assert uses_cuda_stream(func_with_pytorch_stream)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_cupy_stream_creation(self):
+        import cupy
+
+        """Test detection of CuPy CUDA stream creation."""
+
+        def func_with_cupy_stream():
+            stream = cupy.cuda.Stream()
+            return stream
+
+        assert uses_cuda_stream(func_with_cupy_stream)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_generic_stream_creation(self):
+        """Test detection of generic Stream() calls."""
+
+        def func_with_generic_stream():
+            from torch.cuda import Stream
+
+            stream = Stream()
+            return stream
+
+        assert uses_cuda_stream(func_with_generic_stream)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_stream_with_device_id(self):
+        """Test detection of Stream with device ID."""
+
+        def func_with_device_stream():
+            from torch.cuda import Stream
+
+            stream = Stream(0)
+            return stream
+
+        assert uses_cuda_stream(func_with_device_stream)
+
+    def test_no_stream_creation(self):
+        """Test functions without stream creation return False."""
+
+        def func_without_stream():
+            import torch
+
+            x = torch.randn(100, 100)
+            y = x @ x.T
+            return y
+
+        assert not uses_cuda_stream(func_without_stream)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_lambda_function(self):
+        """Test detection in lambda functions."""
+
+        def func_lambda_with_stream():
+            return torch.cuda.Stream()
+
+        def func_lambda_without(x):
+            return x * 2
+
+        assert uses_cuda_stream(func_lambda_with_stream)
+        assert not uses_cuda_stream(func_lambda_without)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_nested_function(self):
+        """Test detection in nested functions."""
+
+        def outer_function():
+            def inner_with_stream():
+                import torch
+
+                return torch.cuda.Stream()
+
+            return inner_with_stream
+
+        inner = outer_function()
+        assert uses_cuda_stream(inner)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_class_method(self):
+        """Test detection in class methods."""
+
+        class StreamClass:
+            def method_with_stream(self):
+                import torch
+
+                self.stream = torch.cuda.Stream()
+                return self.stream
+
+            def method_without_stream(self):
+                return "no stream here"
+
+        obj = StreamClass()
+        assert uses_cuda_stream(obj.method_with_stream)
+        assert not uses_cuda_stream(obj.method_without_stream)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_various_formats(self):
+        """Test various formatting of stream creation."""
+
+        def func_spaces():
+            stream = torch.cuda.Stream()
+            return stream
+
+        def func_multiline():
+            stream = torch.cuda.Stream(device=0)
+            return stream
+
+        def func_chained():
+            result = torch.cuda.Stream().query()
+            return result
+
+        assert uses_cuda_stream(func_spaces)
+        assert uses_cuda_stream(func_multiline)
+        assert uses_cuda_stream(func_chained)
+
+    @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
+    def test_case_sensitivity(self):
+        """Test case-insensitive detection."""
+
+        def func_lowercase():
+            stream = torch.cuda.stream()  # lowercase (if it existed)
+            return stream
+
+        def func_uppercase():
+            stream = torch.cuda.STREAM()  # uppercase (if it existed)
+            return stream
+
+        # These should still be detected due to case-insensitive regex
+        assert uses_cuda_stream(func_lowercase)
+        assert uses_cuda_stream(func_uppercase)
+
+    def test_opoverload_callables(self):
+        """Test that OpOverload objects don't raise exceptions."""
+        import torch
+
+        # Test OpOverload (torch operators)
+        assert not uses_cuda_stream(torch.add)
+        assert not uses_cuda_stream(torch.ops.aten.add)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])