Add Callsite Location Support in TTIR/TTGIR (#190)

Summary:
## Overview

This PR adds support for parsing MLIR callsite locations in Triton IR (TTIR/TTGIR), enabling tritonparse to capture and preserve call stack information for inlined functions.

## Problem

Previously, tritonparse could not parse callsite location definitions like:
```mlir
#loc220 = loc(callsite(#loc57 at #loc190))
```

This caused incomplete source code mappings, losing critical information about:
- Function call chains (which function called which)
- Inlining context (how nested functions are represented)
- Complete source attribution (full path through code to reach an operation)

## Solution

Implemented a **hybrid approach** that:
1. Parses callsite definitions and inherits location from the callee (code being called)
2. Preserves caller references as metadata for call stack traversal
3. Propagates callsite information through the entire mapping pipeline

### Key Design Decisions

- **Callee as primary location**: Maps to the actual code being executed (most relevant for debugging)
- **Metadata preservation**: Stores references (loc IDs) rather than fully expanding call stacks
- **Backward compatible**: Adds optional fields without breaking existing tools
- **Extensible**: Future enhancements can traverse and expand call chains on demand

## Implementation Details

### 1. Added Callsite Pattern (`ir_parser.py`)
```python
CALLSITE_PATTERN = re.compile(
    r"#loc(\d+)\s*=\s*loc\(\s*callsite\(\s*#loc(\d*)\s+at\s+#loc(\d*)\s*\)\s*\)"
)
```

### 2. Enhanced `extract_loc_definitions()` (`ir_parser.py`)
- Collects all callsite definitions during IR parsing
- Resolves callsite references by inheriting location info from callee
- Stores callsite metadata: `is_callsite`, `callsite_callee`, `callsite_caller`
- Validates references with warning messages for undefined locs

### 3. Updated `generate_source_mappings()` (`trace_processor.py`)
- Propagates callsite metadata from `loc_defs` to final `mappings`
- Enables downstream tools to identify and traverse call chains

## Data Structure Example

For a nested callsite like:
```mlir
#loc7 = loc("file.py":1091:8)
#loc57 = loc("file.py":421:16)
#loc58 = loc("file.py":853:16)
#loc190 = loc(callsite(#loc58 at #loc7))
#loc220 = loc(callsite(#loc57 at #loc190))
%0 = tt.load %ptr loc(#loc220)
```

The resulting mapping for line 131 (where `tt.load` is):
```json
{
  "file": "file.py",
  "line": 421,              // From callee (loc57) - actual code executing
  "column": 16,
  "ttir_line": 131,
  "is_callsite": true,
  "callsite_callee": "57",  // Reference to called code
  "callsite_caller": "190"  // Reference to caller (can traverse chain)
}
```

**Call chain represented**: `_ragged_hstu_attn_fwd` (1091:8) → `_ragged_hstu_attn_fwd_compute` (853:16) → `_ragged_hstu_attn_fwd_one_block` (421:16) ← **executing here**

## Testing

Added comprehensive unit tests in `tests/test_tritonparse.py`:
- `TestTritonparseCPU::test_callsite_parsing`
- Validates nested callsite parsing
- Verifies metadata propagation to mappings
- Tests both simple and nested callsite scenarios

**All tests pass ✅**

## Impact

### Benefits
- ✅ Complete source mapping for inlined functions
- ✅ Preserves call stack information for debugging
- ✅ Enables future call chain visualization
- ✅ Backward compatible with existing tools

### No Breaking Changes
- Only adds optional fields to existing mappings
- Existing code that doesn't check for callsite fields continues to work
- No changes to public APIs

## Files Changed

1. `tritonparse/ir_parser.py` - Added callsite parsing logic
2. `tritonparse/trace_processor.py` - Propagate callsite metadata
3. `tests/test_tritonparse.py` - Added unit tests
4. `CALLSITE_IMPLEMENTATION.md` - Detailed implementation documentation

## Future Work

Potential enhancements (not in this PR):
1. Automatic call stack expansion utility functions
2. Call stack caching for performance optimization
3. Frontend UI support for call chain visualization
4. Cycle detection for complex callsite graphs

Pull Request resolved: #190

Reviewed By: njriasan

Differential Revision: D85681542

Pulled By: FindHao

fbshipit-source-id: aeadd4cb65a453f50a8d48ca7ea68a40515a1f07

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

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,50 @@ def resolve_alias(current_id: str) -> Dict[str, Any]:
     for alias_id, target_id in alias_map.items():
         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:
+            locations[k]["def_line"] = v
+    for alias_id, target_id in alias_map.items():
         if alias_id in locations:
             locations[alias_id]["alias_of"] = target_id
             if alias_id in alias_name_map:

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"]