Add location alias parsing support in IR parser and trace processor

Implements comprehensive support for parsing location aliases in TTIR/TTGIR:

Backend changes:
- Add ALIAS_WITH_NAME_PATTERN and ALIAS_SIMPLE_PATTERN regex to match
  alias definitions like #loc16 = loc("pid"(#loc2)) and #loc20 = loc(#loc16)
- Update LOC_PATTERN to only match numbered locs (#loc1, #loc2) not bare #loc
- Implement alias resolution logic with cycle detection
- Track definition lines (def_line), alias names (alias_name), and
  alias targets (alias_of) for each location
- Propagate alias metadata through trace_processor to frontend
- Add separate mappings for loc definition lines with kind="loc_def"

Testing:
- Add comprehensive unit test test_loc_alias_parsing() covering:
  * Bare #loc references (e.g., #loc13 = loc("x_ptr"(#loc)))
  * Named aliases with numbered refs (e.g., #loc16 = loc("pid"(#loc2)))
  * Simple aliases without names (e.g., #loc20 = loc(#loc1))
  * Definition line tracking for all loc types

This enables the frontend to display alias information and highlight
loc definition lines in the code viewer.

Author: FindHao

tests/test_tritonparse.py

@@ -1442,6 +1442,70 @@ def count_all_blobs(manager_dir_path):
             # Cleanup test-specific cache
             self.cleanup_test_cache(test_cache_dir, prev_cache_dir)
 
+    def test_loc_alias_parsing(self):
+        """Test parsing of location aliases in TTIR/TTGIR"""
+        from tritonparse.ir_parser import extract_loc_definitions
+
+        # Test case 1: Bare #loc reference (no number)
+        ir_with_bare_loc = """
+module {
+  #loc = loc("/tmp/test.py":10:5)
+  #loc13 = loc("x_ptr"(#loc))
+  func @kernel(%arg0: !tt.ptr<f32> loc(#loc13)) {
+    return loc(#loc)
+  }
+}
+"""
+        locs = extract_loc_definitions(ir_with_bare_loc)
+        # Main #loc should be stored with "" key
+        assert "" in locs, "Main #loc not found"
+        assert locs[""]["file"] == "/tmp/test.py"
+        assert locs[""]["line"] == 10
+        # Alias #loc13 should resolve to same location
+        assert "13" in locs, "#loc13 not found"
+        assert locs["13"]["file"] == "/tmp/test.py"
+        assert locs["13"]["line"] == 10
+        assert locs["13"]["alias_name"] == "x_ptr"
+        assert locs["13"]["alias_of"] == ""
+
+        # Test case 2: Named alias with numbered reference
+        ir_with_numbered_alias = """
+#loc = loc("/tmp/test.py":5:0)
+#loc2 = loc("/tmp/test.py":20:28)
+#loc16 = loc("pid"(#loc2))
+%0 = tt.get_program_id x : i32 loc(#loc16)
+"""
+        locs = extract_loc_definitions(ir_with_numbered_alias)
+        assert "2" in locs
+        assert locs["2"]["line"] == 20
+        assert "16" in locs
+        assert locs["16"]["file"] == "/tmp/test.py"
+        assert locs["16"]["line"] == 20
+        assert locs["16"]["alias_name"] == "pid"
+        assert locs["16"]["alias_of"] == "2"
+
+        # Test case 3: Simple alias (no name)
+        ir_with_simple_alias = """
+#loc = loc("/tmp/test.py":1:1)
+#loc1 = loc("/tmp/test.py":15:10)
+#loc20 = loc(#loc1)
+%1 = arith.constant 0 : i32 loc(#loc20)
+"""
+        locs = extract_loc_definitions(ir_with_simple_alias)
+        assert "1" in locs
+        assert "20" in locs
+        assert locs["20"]["file"] == "/tmp/test.py"
+        assert locs["20"]["line"] == 15
+        assert locs["20"]["alias_of"] == "1"
+        assert "alias_name" not in locs["20"]
+
+        # Test case 4: Definition line tracking
+        assert "def_line" in locs[""]
+        assert "def_line" in locs["1"]
+        assert "def_line" in locs["20"]
+
+        print("✓ All loc alias parsing tests passed")
+
 
 if __name__ == "__main__":
     unittest.main()

tritonparse/ir_parser.py

@@ -10,7 +10,8 @@
 
 # the definition of the #loc directive. they are in the bottom of the IR files
 # Example:#loc2 = loc("/tmp/torchinductor_yhao/yp/abcdef.py":20:28)
-LOC_PATTERN = re.compile(r'#loc(\d*) = loc\("([^"]+)":(\d+):(\d+)\)')
+# Note: This should only match numbered locs like #loc1, #loc2, not bare #loc
+LOC_PATTERN = re.compile(r'#loc(\d+) = loc\("([^"]+)":(\d+):(\d+)\)')
 
 # the reference to the #loc directive. they are in the end of lines of the IR files
 # Example: loc(#loc2)
@@ -33,6 +34,17 @@
 )
 
 
+# alias loc definitions in TTGIR/TTIR
+# Example: #loc16 = loc("pid"(#loc2))
+# Example: #loc13 = loc("x_ptr"(#loc)) - bare #loc without number
+ALIAS_WITH_NAME_PATTERN = re.compile(
+    r'#loc(\d+)\s*=\s*loc\("([^"]+)"\s*\(\s*#loc(\d*)\s*\)\s*\)'
+)
+
+# Example: #loc20 = loc(#loc16)
+ALIAS_SIMPLE_PATTERN = re.compile(r"#loc(\d+)\s*=\s*loc\(\s*#loc(\d*)\s*\)")
+
+
 def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     """
     Extracts location definitions from the given IR content.
@@ -50,6 +62,7 @@ def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     """
     locations = {}
     # The first #loc directive is a special case. It locates at the top of the IR files
+    # Store it with empty string "" as key to avoid conflict with #loc1
     main_match = re.search(r'#loc = loc\("([^"]+)":(\d+):(\d+)\)', ir_content)
     if main_match:
         locations[""] = {
@@ -61,6 +74,84 @@ def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     for loc_id, filename, line, col in LOC_PATTERN.findall(ir_content):
         key = loc_id
         locations[key] = {"file": filename, "line": int(line), "column": int(col)}
+
+    # Handle alias-style loc definitions that reference another #loc
+    # Build alias map first: alias_id -> target_id
+    alias_map: Dict[str, str] = {}
+    for m in ALIAS_WITH_NAME_PATTERN.finditer(ir_content):
+        alias_id, _name, target_id = m.groups()
+        # Empty target_id means bare #loc, map to "" (main loc key)
+        alias_map[alias_id] = target_id or ""
+    for m in ALIAS_SIMPLE_PATTERN.finditer(ir_content):
+        alias_id, target_id = m.groups()
+        # Empty target_id means bare #loc, map to "" (main loc key)
+        alias_map[alias_id] = target_id or ""
+
+    # Build definition line map and alias name map by scanning lines
+    def_line_map: Dict[str, int] = {}
+    alias_name_map: Dict[str, str] = {}
+    main_loc_line: int = 0
+    for i, line in enumerate(ir_content.split("\n"), start=1):
+        if m := ALIAS_WITH_NAME_PATTERN.search(line):
+            alias_id, name, target_id = m.groups()
+            def_line_map[alias_id] = i
+            alias_name_map[alias_id] = name
+            # ensure alias map is populated even if only found in line scan
+            # Empty target_id means bare #loc, map to "" (main loc key)
+            alias_map.setdefault(alias_id, target_id or "")
+        elif m := ALIAS_SIMPLE_PATTERN.search(line):
+            alias_id, target_id = m.groups()
+            def_line_map[alias_id] = i
+            # Empty target_id means bare #loc, map to "" (main loc key)
+            alias_map.setdefault(alias_id, target_id or "")
+        if m2 := LOC_PATTERN.search(line):
+            base_id, _fn, _ln, _col = m2.groups()
+            def_line_map[base_id] = i
+        if re.search(r'#loc\s*=\s*loc\("[^"]+":\d+:\d+\)', line):
+            # main #loc = loc("file":line:col) without id
+            main_loc_line = main_loc_line or i
+
+    # Resolve aliases to base locations (file/line/column)
+    resolving_stack = set()
+
+    def resolve_alias(current_id: str) -> Dict[str, Any]:
+        # Already a concrete location
+        if current_id in locations:
+            return locations[current_id]
+        # Detect cycles
+        if current_id in resolving_stack:
+            return {}
+        resolving_stack.add(current_id)
+        parent_id = alias_map.get(current_id)
+        result: Dict[str, Any] = {}
+        if parent_id is not None:
+            base = resolve_alias(parent_id)
+            if base:
+                # copy to avoid sharing the same dict by reference
+                result = {
+                    "file": base.get("file"),
+                    "line": base.get("line"),
+                    "column": base.get("column"),
+                }
+                locations[current_id] = result
+        resolving_stack.remove(current_id)
+        return result
+
+    for alias_id in list(alias_map.keys()):
+        if alias_id not in locations:
+            resolve_alias(alias_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:
+                locations[alias_id]["alias_name"] = alias_name_map[alias_id]
+    if main_loc_line and "" in locations:
+        locations[""]["def_line"] = main_loc_line
     return locations
 
 

tritonparse/trace_processor.py

@@ -71,12 +71,37 @@ def generate_source_mappings(
             }
         elif loc_id in loc_defs:
             info = loc_defs[loc_id]
-            mappings[str(ln)] = {
+            entry = {
                 "file": info["file"],
                 "line": info["line"],
                 "column": info["column"],
                 f"{ir_type}_line": ln,
             }
+            # Propagate alias metadata if present
+            if "alias_name" in info:
+                entry["alias_name"] = info["alias_name"]
+            if "alias_of" in info:
+                entry["loc_id"] = loc_id
+            mappings[str(ln)] = entry
+
+    # Add separate entries for loc definition lines
+    for loc_id, info in loc_defs.items():
+        if "def_line" in info:
+            def_ln = info["def_line"]
+            # Only create mapping if this line doesn't already have one
+            if str(def_ln) not in mappings:
+                entry = {
+                    "file": info["file"],
+                    "line": info["line"],
+                    "column": info["column"],
+                    f"{ir_type}_line": def_ln,
+                    "kind": "loc_def",
+                }
+                if "alias_name" in info:
+                    entry["alias_name"] = info["alias_name"]
+                if "alias_of" in info:
+                    entry["loc_id"] = loc_id
+                mappings[str(def_ln)] = entry
 
     return mappings