tritonparse 0.3.2.dev20251210071601__py3-none-any.whl

tritonparse/reproducer/templates/utils.py

@@ -0,0 +1,48 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+import inspect
+import logging
+from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+
+def get_function_source(
+    func: Callable[[], None], with_invocation: bool = True
+) -> list[str]:
+    """
+    Extract function source code and optionally include invocation.
+
+    Args:
+        func: Function to extract source code from
+        with_invocation: Whether to include function invocation code
+
+    Returns:
+        List containing source code and optional invocation statement
+    """
+    source = inspect.getsource(func).rstrip()
+    result = [source]
+
+    if with_invocation:
+        result.append("")
+        result.append(f"{func.__name__}()")
+
+    return result
+
+
+def _disable_triton_autotune() -> None:
+    """
+    Monkey patch the triton.autotune decorator to skip autotuning entirely.
+    """
+    logger.info("Disabling triton autotune")
+
+    def dummy_autotune(configs, key=None, **kwargs):
+        def decorator(func):
+            return func  # Just pass through, let @triton.jit handle compilation
+
+        return decorator
+
+    import triton
+
+    triton.autotune = dummy_autotune
+    logger.info("Disabled triton autotune")

tritonparse/reproducer/tests/artifacts/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+"""Test artifacts for multi-file call graph analyzer integration tests."""

tritonparse/reproducer/tests/artifacts/triton_fused_kernel.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+if __name__ == "__main__":
+    import sys
+    from pathlib import Path
+
+    # Add tritonparse_root to sys.path for module resolution
+    tritonparse_root = Path(__file__).resolve().parents[4]
+    if str(tritonparse_root) not in sys.path:
+        sys.path.insert(0, str(tritonparse_root))
+
+import torch
+import triton
+import triton.language as tl
+from tritonparse.reproducer.tests.artifacts.triton_preprocess import scale_kernel
+
+
+@triton.jit
+def main_kernel(
+    input_ptr,
+    output_ptr,
+    n_elements,
+    scale: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+    scaled = scale_kernel(x, scale)
+    result = scaled * 2.0
+    tl.store(output_ptr + offsets, result, mask=mask)
+
+
+def launch_main_kernel() -> None:
+    """Launch and test the main_kernel on GPU."""
+    if not torch.cuda.is_available():
+        print("CUDA not available - showing call graph only")
+        print("  main_kernel -> scale_kernel -> add_values")
+        return
+
+    size = 1024
+    scale_factor = 3.0
+    BLOCK_SIZE = 256
+
+    x = torch.randn(size, device="cuda", dtype=torch.float32)
+    output = torch.zeros_like(x)
+
+    main_kernel[(256,)](
+        input_ptr=x,
+        output_ptr=output,
+        n_elements=size,
+        scale=scale_factor,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+    print("✅ Kernel executed successfully")
+
+
+if __name__ == "__main__":
+    launch_main_kernel()

tritonparse/reproducer/tests/artifacts/triton_preprocess.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+import triton
+import triton.language as tl
+from tritonparse.reproducer.tests.artifacts.triton_utils import add_values
+
+
+@triton.jit
+def scale_kernel(
+    x: tl.tensor,
+    scale: tl.constexpr,
+) -> tl.tensor:
+    result = x * scale
+    return add_values(result, 0.0)

tritonparse/reproducer/tests/artifacts/triton_utils.py

@@ -0,0 +1,14 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def add_values(
+    a: tl.tensor,
+    b: tl.tensor,
+) -> tl.tensor:
+    return a + b

tritonparse/reproducer/tests/test_import_parser.py

@@ -0,0 +1,164 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+import ast
+import unittest
+from pathlib import Path
+
+from tritonparse.reproducer.import_parser import ImportParser
+from tritonparse.reproducer.import_resolver import ImportResolver
+
+
+class TestImportParser(unittest.TestCase):
+    """Test ImportParser functionality."""
+
+    def setUp(self) -> None:
+        """Set up test fixtures."""
+        test_dir = Path(__file__).resolve().parent
+        reproducer_dir = test_dir.parent
+        self.project_root = str(reproducer_dir.parent)
+        self.resolver = ImportResolver(self.project_root)
+        self.parser = ImportParser(self.resolver)
+
+    def test_parse_simple_import(self) -> None:
+        """Test parsing simple 'import X' statement."""
+        code = """
+import os
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 1)
+        self.assertEqual(imports[0].import_type, "import")
+        self.assertEqual(imports[0].module, "os")
+        self.assertEqual(imports[0].names, ["os"])
+        self.assertTrue(imports[0].is_external)
+
+    def test_parse_import_with_alias(self) -> None:
+        """Test parsing 'import X as Y' statement."""
+        code = """
+import numpy as np
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 1)
+        self.assertEqual(imports[0].module, "numpy")
+        self.assertIn("np", imports[0].aliases)
+        self.assertEqual(imports[0].aliases["np"], "numpy")
+
+    def test_parse_from_import(self) -> None:
+        """Test parsing 'from X import Y' statement."""
+        code = """
+from typing import List, Dict
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 2)
+        self.assertEqual(imports[0].import_type, "from_import")
+        self.assertEqual(imports[0].module, "typing")
+        self.assertIn("List", imports[0].names)
+        self.assertIn("Dict", imports[1].names)
+
+    def test_parse_from_import_with_alias(self) -> None:
+        """Test parsing 'from X import Y as Z' statement."""
+        code = """
+from collections import OrderedDict as OD
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 1)
+        self.assertEqual(imports[0].names, ["OrderedDict"])
+        self.assertIn("OD", imports[0].aliases)
+        self.assertEqual(imports[0].aliases["OD"], "OrderedDict")
+
+    def test_parse_relative_import_level_1(self) -> None:
+        """Test parsing relative import 'from . import X'.
+
+        When level=1 (from . import X), we're importing from the current package.
+        With package="pytorch.tritonparse.reproducer", level=1 means current package,
+        so the full module is "pytorch.tritonparse.reproducer".
+        """
+        code = """
+from . import utils
+"""
+        tree = ast.parse(code, filename="test.py")
+        package = "tritonparse.reproducer"
+        imports = self.parser.parse_imports(tree, "test.py", package)
+
+        self.assertEqual(len(imports), 1)
+        self.assertEqual(imports[0].level, 1)
+        self.assertEqual(imports[0].module, "tritonparse.reproducer")
+        self.assertEqual(imports[0].names, ["utils"])
+
+    def test_parse_relative_import_level_2(self) -> None:
+        """Test parsing relative import 'from .. import X'.
+
+        When level=2 (from ..utils import X), we're importing from the parent package.
+        With package="pytorch.tritonparse.reproducer.submodule", level=2 removes 1 component,
+        giving us "pytorch.tritonparse.reproducer", then we append "utils" to get
+        "pytorch.tritonparse.reproducer.utils".
+        """
+        code = """
+from ..utils import helper
+"""
+        tree = ast.parse(code, filename="test.py")
+        package = "tritonparse.reproducer.submodule"
+        imports = self.parser.parse_imports(tree, "test.py", package)
+
+        self.assertEqual(len(imports), 1)
+        self.assertEqual(imports[0].level, 2)
+        self.assertEqual(imports[0].module, "tritonparse.reproducer.utils")
+        self.assertEqual(imports[0].names, ["helper"])
+
+    def test_parse_multiple_imports(self) -> None:
+        """Test parsing multiple import statements."""
+        code = """
+import os
+import sys
+from typing import List
+from collections import defaultdict
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 4)
+        self.assertEqual(imports[0].module, "os")
+        self.assertEqual(imports[1].module, "sys")
+        self.assertEqual(imports[2].module, "typing")
+        self.assertEqual(imports[3].module, "collections")
+
+    def test_parse_project_internal_import(self) -> None:
+        """Test parsing imports from within the project."""
+        code = """
+from tritonparse.reproducer import ast_analyzer
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 1)
+        self.assertFalse(imports[0].is_external)
+        self.assertIsNotNone(imports[0].resolved_path)
+        if imports[0].resolved_path:
+            self.assertTrue(imports[0].resolved_path.startswith(self.project_root))
+
+    def test_parse_lineno_tracking(self) -> None:
+        """Test that line numbers are correctly tracked."""
+        code = """
+import os
+
+from typing import List
+
+
+import sys
+"""
+        tree = ast.parse(code, filename="test.py")
+        imports = self.parser.parse_imports(tree, "test.py")
+
+        self.assertEqual(len(imports), 3)
+        self.assertEqual(imports[0].lineno, 2)  # import os
+        self.assertEqual(imports[1].lineno, 4)  # from typing import List
+        self.assertEqual(imports[2].lineno, 7)  # import sys

tritonparse/reproducer/tests/test_import_resolver.py

@@ -0,0 +1,88 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Unit tests for ImportResolver.
+
+Tests the ImportResolver implementation to ensure it correctly:
+1. Resolves internal imports
+2. Detects external modules
+3. Handles non-existent modules gracefully
+"""
+
+import unittest
+from pathlib import Path
+
+from tritonparse.reproducer.import_resolver import ImportResolver
+
+
+class ImportResolverTest(unittest.TestCase):
+    """Unit tests for ImportResolver class."""
+
+    def setUp(self) -> None:
+        """Set up test fixtures."""
+        # Use the tritonparse directory as project root for testing
+        test_dir = Path(__file__).resolve().parent
+        reproducer_dir = test_dir.parent
+        self.project_root = str(reproducer_dir.parent)
+        self.resolver = ImportResolver(project_root=self.project_root)
+
+    def test_resolve_tritonparse_module(self) -> None:
+        """Test resolving a module within tritonparse."""
+        # Setup: module that exists in tritonparse
+        module_name = "reproducer.ast_analyzer"
+
+        # Execute: resolve the module
+        path, is_external = self.resolver.resolve_import(module_name)
+
+        # Assert: should resolve to tritonparse path
+        self.assertIsNotNone(path)
+        assert path is not None  # For type checker
+        self.assertTrue(path.startswith(self.project_root))
+        self.assertTrue(Path(path).exists())
+        self.assertFalse(is_external)
+
+    def test_external_module_torch(self) -> None:
+        """Test that torch is correctly identified as external."""
+        # Setup: torch is a known external module
+        module_name = "torch"
+
+        # Execute: resolve the module
+        path, is_external = self.resolver.resolve_import(module_name)
+
+        # Assert: should be external
+        self.assertIsNone(path)
+        self.assertTrue(is_external)
+
+    def test_nonexistent_module(self) -> None:
+        """Test handling of non-existent module."""
+        # Setup: a module that doesn't exist
+        module_name = "this_module_does_not_exist_anywhere"
+
+        # Execute: resolve the module
+        path, is_external = self.resolver.resolve_import(module_name)
+
+        # Assert: should handle gracefully as external
+        self.assertIsNone(path)
+        self.assertTrue(is_external)
+
+    def test_is_external_module_torch(self) -> None:
+        """Test is_external_module() for torch."""
+        # Setup: torch module name
+        module_name = "torch"
+
+        # Execute: check if external
+        result = self.resolver.is_external_module(module_name)
+
+        # Assert: should be true
+        self.assertTrue(result)
+
+    def test_is_external_module_internal(self) -> None:
+        """Test is_external_module() for internal module."""
+        # Setup: internal module name
+        module_name = "reproducer.ast_analyzer"
+
+        # Execute: check if external
+        result = self.resolver.is_external_module(module_name)
+
+        # Assert: should be false
+        self.assertFalse(result)

tritonparse/reproducer/tests/test_multi_file_analyzer.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+"""
+Tests for MultiFileCallGraphAnalyzer.
+
+This test suite validates the multi-file call graph analysis functionality,
+including cross-file dependency tracking, import resolution, and result consolidation.
+"""
+
+import unittest
+from pathlib import Path
+
+from tritonparse.reproducer.multi_file_analyzer import MultiFileCallGraphAnalyzer
+
+
+class TestMultiFileCallGraphAnalyzer(unittest.TestCase):
+    """Test the MultiFileCallGraphAnalyzer with real Triton test files."""
+
+    def setUp(self) -> None:
+        """Set up test fixtures."""
+        # Get code root (4 levels up from this test file)
+        test_file = Path(__file__).resolve()
+        self.code_root = str(test_file.parents[4])
+
+        # Test artifacts directory
+        self.artifacts_dir = test_file.parent / "artifacts"
+        self.entry_file = str(self.artifacts_dir / "triton_fused_kernel.py")
+        self.preprocess_file = str(self.artifacts_dir / "triton_preprocess.py")
+        self.utils_file = str(self.artifacts_dir / "triton_utils.py")
+
+    def test_single_file_analysis(self) -> None:
+        """Test analysis of a single file without following imports."""
+        # Setup: Analyze only triton_utils.py (which has no imports to internal files)
+        analyzer = MultiFileCallGraphAnalyzer(
+            entry_file=self.utils_file,
+            entry_function="add_values",
+        )
+
+        # Execute: Run analysis
+        result = analyzer.analyze()
+
+        # Assert: Only one file analyzed (triton_utils.py)
+        self.assertEqual(result.stats.total_files_analyzed, 1)
+        self.assertIn(self.utils_file, result.analyzed_files)
+
+        # Assert: add_values is the backend, should not be in dependent functions
+        # (dependent functions excludes the backend itself)
+        self.assertNotIn(
+            "tritonparse.reproducer.tests.artifacts.triton_utils.add_values",
+            result.functions,
+        )
+
+        # Assert: All imports should be external (triton, tl)
+        self.assertTrue(all(imp.is_external for imp in result.imports))
+
+    def test_multi_file_traversal(self) -> None:
+        """Test that analyzer follows imports across multiple files."""
+        # Setup: Start from main_kernel in triton_fused_kernel.py
+        analyzer = MultiFileCallGraphAnalyzer(
+            entry_file=self.entry_file,
+            entry_function="main_kernel",
+        )
+
+        # Execute: Run analysis
+        result = analyzer.analyze()
+
+        # Assert: Should analyze all 3 files in the call chain
+        # main_kernel -> scale_kernel -> add_values
+        self.assertEqual(result.stats.total_files_analyzed, 3)
+
+        # Assert: All 3 kernel files should be included
+        self.assertIn(self.entry_file, result.analyzed_files)
+        self.assertIn(self.preprocess_file, result.analyzed_files)
+        self.assertIn(self.utils_file, result.analyzed_files)
+
+        # Assert: Verify scale_kernel and add_values are in result.functions
+        function_names = set(result.functions.keys())
+        self.assertEqual(len(result.functions.items()), 2)
+        self.assertTrue(
+            any("scale_kernel" in name for name in function_names),
+            f"scale_kernel should be in functions. Found: {function_names}",
+        )
+        self.assertTrue(
+            any("add_values" in name for name in function_names),
+            f"add_values should be in functions. Found: {function_names}",
+        )
+
+        # Assert: Verify source code was properly extracted with exact lengths
+        # Find the full qualified names for these functions
+        scale_kernel_name = next(
+            name for name in function_names if "scale_kernel" in name
+        )
+        add_values_name = next(name for name in function_names if "add_values" in name)
+
+        # Verify source code has exact expected lengths
+        scale_kernel_source = result.functions[scale_kernel_name]
+        add_values_source = result.functions[add_values_name]
+
+        self.assertGreaterEqual(
+            len(scale_kernel_source),
+            50,
+            f"scale_kernel source code should be exactly 50 chars. Got: {len(scale_kernel_source)} chars",
+        )
+        self.assertGreaterEqual(
+            len(add_values_source),
+            50,
+            f"add_values source code should be exactly 50 chars. Got: {len(add_values_source)} chars",
+        )
+
+        # Verify short names are correctly extracted
+        self.assertEqual(result.function_short_names[scale_kernel_name], "scale_kernel")
+        self.assertEqual(result.function_short_names[add_values_name], "add_values")
+
+
+if __name__ == "__main__":
+    unittest.main()

tritonparse/reproducer/types.py

@@ -0,0 +1,20 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from enum import Enum
+
+
+class KernelImportMode(str, Enum):
+    """
+    Kernel import strategy for reproducer generation.
+
+    Inherits from str to allow direct string comparison and use in argparse.
+
+    Attributes:
+        DEFAULT: Import kernel from original file (current behavior).
+        COPY: Embed kernel source code directly in reproducer.
+        OVERRIDE_TTIR: Use TTIR from compilation event with monkeypatch.
+    """
+
+    DEFAULT = "default"
+    COPY = "copy"
+    OVERRIDE_TTIR = "override-ttir"