Add MLIR callsite location parsing support

Implement parsing for MLIR callsite locations (e.g., `loc(callsite(#loc57 at #loc190))`)
to preserve function call chain information for inlined functions.

Changes:
- Add CALLSITE_PATTERN regex to recognize callsite definitions
- Extend extract_loc_definitions() to collect and resolve callsite references
- Update generate_source_mappings() to propagate callsite metadata
- Add comprehensive unit tests in TestTritonparseCPU::test_callsite_parsing

Callsite locations inherit file/line/column from the callee (actual code executing)
while preserving caller references as metadata for call stack traversal. This enables
complete source mapping for inlined functions without breaking existing tools.

Test Plan:
python -m pytest tests/test_tritonparse.py::TestTritonparseCPU::test_callsite_parsing -v

Author: FindHao

tests/test_tritonparse.py

@@ -139,6 +139,63 @@ def clear_all_caches(*kernels):
 class TestTritonparseCPU(unittest.TestCase):
     """CPU-only tests (no CUDA required)"""
 
+    def test_callsite_parsing(self):
+        """Test parsing of callsite locations in TTIR/TTGIR"""
+        from tritonparse.ir_parser import extract_loc_definitions
+        from tritonparse.trace_processor import generate_source_mappings
+
+        # Test MLIR callsite location definitions
+        ir_with_callsite = """
+module {
+  #loc7 = loc("/tmp/test.py":1091:8)
+  #loc57 = loc("/tmp/test.py":421:16)
+  #loc58 = loc("/tmp/test.py":853:16)
+  #loc190 = loc(callsite(#loc58 at #loc7))
+  #loc220 = loc(callsite(#loc57 at #loc190))
+  %0 = tt.load %ptr loc(#loc220)
+}
+"""
+        # Extract loc definitions
+        locs = extract_loc_definitions(ir_with_callsite)
+
+        # Verify loc220 (nested callsite)
+        self.assertIn("220", locs)
+        self.assertEqual(locs["220"]["file"], "/tmp/test.py")
+        self.assertEqual(locs["220"]["line"], 421)  # Inherited from callee loc57
+        self.assertEqual(locs["220"]["column"], 16)
+        self.assertTrue(locs["220"].get("is_callsite"))
+        self.assertEqual(locs["220"]["callsite_callee"], "57")
+        self.assertEqual(locs["220"]["callsite_caller"], "190")
+
+        # Verify loc190 (simple callsite)
+        self.assertIn("190", locs)
+        self.assertEqual(locs["190"]["line"], 853)  # Inherited from callee loc58
+        self.assertTrue(locs["190"].get("is_callsite"))
+        self.assertEqual(locs["190"]["callsite_callee"], "58")
+        self.assertEqual(locs["190"]["callsite_caller"], "7")
+
+        # Test source mappings generation
+        mappings = generate_source_mappings(ir_with_callsite, "ttir")
+
+        # Find the line with tt.load
+        line_with_load = None
+        for line_num, content in enumerate(ir_with_callsite.split("\n"), start=1):
+            if "tt.load" in content:
+                line_with_load = str(line_num)
+                break
+
+        self.assertIsNotNone(line_with_load)
+        self.assertIn(line_with_load, mappings)
+
+        mapping = mappings[line_with_load]
+        self.assertEqual(mapping["file"], "/tmp/test.py")
+        self.assertEqual(mapping["line"], 421)  # From loc220 -> loc57
+        self.assertTrue(mapping.get("is_callsite"))
+        self.assertEqual(mapping["callsite_callee"], "57")
+        self.assertEqual(mapping["callsite_caller"], "190")
+
+        print("✓ Callsite parsing tests passed")
+
     def test_convert(self):
         """Test convert function with various data types"""
         # Test with primitive types
@@ -1506,6 +1563,63 @@ def test_loc_alias_parsing(self):
 
         print("✓ All loc alias parsing tests passed")
 
+    def test_callsite_parsing(self):
+        """Test parsing of callsite locations in TTIR/TTGIR"""
+        from tritonparse.ir_parser import extract_loc_definitions
+        from tritonparse.trace_processor import generate_source_mappings
+
+        # Test MLIR callsite location definitions
+        ir_with_callsite = """
+module {
+  #loc7 = loc("/tmp/test.py":1091:8)
+  #loc57 = loc("/tmp/test.py":421:16)
+  #loc58 = loc("/tmp/test.py":853:16)
+  #loc190 = loc(callsite(#loc58 at #loc7))
+  #loc220 = loc(callsite(#loc57 at #loc190))
+  %0 = tt.load %ptr loc(#loc220)
+}
+"""
+        # Extract loc definitions
+        locs = extract_loc_definitions(ir_with_callsite)
+
+        # Verify loc220 (nested callsite)
+        self.assertIn("220", locs)
+        self.assertEqual(locs["220"]["file"], "/tmp/test.py")
+        self.assertEqual(locs["220"]["line"], 421)  # Inherited from callee loc57
+        self.assertEqual(locs["220"]["column"], 16)
+        self.assertTrue(locs["220"].get("is_callsite"))
+        self.assertEqual(locs["220"]["callsite_callee"], "57")
+        self.assertEqual(locs["220"]["callsite_caller"], "190")
+
+        # Verify loc190 (simple callsite)
+        self.assertIn("190", locs)
+        self.assertEqual(locs["190"]["line"], 853)  # Inherited from callee loc58
+        self.assertTrue(locs["190"].get("is_callsite"))
+        self.assertEqual(locs["190"]["callsite_callee"], "58")
+        self.assertEqual(locs["190"]["callsite_caller"], "7")
+
+        # Test source mappings generation
+        mappings = generate_source_mappings(ir_with_callsite, "ttir")
+
+        # Find the line with tt.load
+        line_with_load = None
+        for line_num, content in enumerate(ir_with_callsite.split("\n"), start=1):
+            if "tt.load" in content:
+                line_with_load = str(line_num)
+                break
+
+        self.assertIsNotNone(line_with_load)
+        self.assertIn(line_with_load, mappings)
+
+        mapping = mappings[line_with_load]
+        self.assertEqual(mapping["file"], "/tmp/test.py")
+        self.assertEqual(mapping["line"], 421)  # From loc220 -> loc57
+        self.assertTrue(mapping.get("is_callsite"))
+        self.assertEqual(mapping["callsite_callee"], "57")
+        self.assertEqual(mapping["callsite_caller"], "190")
+
+        print("✓ Callsite parsing tests passed")
+
 
 if __name__ == "__main__":
     unittest.main()

tritonparse/ir_parser.py

@@ -44,6 +44,14 @@
 # Example: #loc20 = loc(#loc16)
 ALIAS_SIMPLE_PATTERN = re.compile(r"#loc(\d+)\s*=\s*loc\(\s*#loc(\d*)\s*\)")
 
+# Callsite loc definitions in TTIR/TTGIR
+# Example: #loc220 = loc(callsite(#loc57 at #loc190))
+# Captures: loc_id, callee_loc_id, caller_loc_id
+# Note: Uses (\d*) to match optional numbers (for bare #loc references)
+CALLSITE_PATTERN = re.compile(
+    r"#loc(\d+)\s*=\s*loc\(\s*callsite\(\s*#loc(\d*)\s+at\s+#loc(\d*)\s*\)\s*\)"
+)
+
 
 def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     """
@@ -141,6 +149,44 @@ def resolve_alias(current_id: str) -> Dict[str, Any]:
         if alias_id not in locations:
             resolve_alias(alias_id)
 
+    # Collect callsite definitions
+    callsite_defs = []
+    for i, line in enumerate(ir_content.split("\n"), start=1):
+        if m := CALLSITE_PATTERN.search(line):
+            loc_id, callee_id, caller_id = m.groups()
+            # Empty strings map to main loc key ""
+            callsite_defs.append((loc_id, callee_id or "", caller_id or "", i))
+
+    # Resolve callsite definitions
+    # A callsite inherits the location from its callee (the code being called)
+    # and stores a reference to its caller (the code doing the calling)
+    for loc_id, callee_id, caller_id, def_line in callsite_defs:
+        if loc_id not in locations:  # Avoid overwriting existing definitions
+            if callee_id in locations:
+                # Inherit location info from callee
+                callee_info = locations[callee_id]
+                locations[loc_id] = {
+                    "file": callee_info["file"],
+                    "line": callee_info["line"],
+                    "column": callee_info["column"],
+                    "def_line": def_line,
+                    "is_callsite": True,
+                    "callsite_callee": callee_id,
+                    "callsite_caller": caller_id,
+                }
+            else:
+                logger.warning(
+                    f"Callsite #loc{loc_id} references undefined callee #loc{callee_id}"
+                )
+                # Note: We don't add this callsite to locations since callee is missing
+
+    # Verify caller references (warning only, don't block)
+    for loc_id, _callee_id, caller_id, _def_line in callsite_defs:
+        if loc_id in locations and caller_id and caller_id not in locations:
+            logger.warning(
+                f"Callsite #loc{loc_id} references undefined caller #loc{caller_id}"
+            )
+
     # Attach definition line and alias metadata
     for k, v in def_line_map.items():
         if k in locations:

tritonparse/trace_processor.py

@@ -77,6 +77,11 @@ def generate_source_mappings(
                 "column": info["column"],
                 f"{ir_type}_line": ln,
             }
+            # Propagate callsite metadata if present
+            if info.get("is_callsite"):
+                entry["is_callsite"] = True
+                entry["callsite_callee"] = info["callsite_callee"]
+                entry["callsite_caller"] = info["callsite_caller"]
             # Propagate alias metadata if present
             if "alias_name" in info:
                 entry["alias_name"] = info["alias_name"]