tritonparse 0.3.1.dev20251029071541__py3-none-any.whl → 0.3.1.dev20251031071507__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of tritonparse might be problematic. Click here for more details.

@@ -46,20 +46,349 @@ def process_amd_gcn_bufferops(
46
46
  return process_amd_bufferop(ir_content, io_keys)
47
47
 
48
48
 
49
+ def find_loop_bounds(ir_content: str) -> list[tuple[int, int]]:
50
+ """
51
+ Find the bounds of all scf.for loops in the IR content.
52
+ These are the only candidates for Software Pipelining (SWP).
53
+
54
+ A loop starts with 'scf.for' and ends when its closing brace '}' is found.
55
+ Brace counts are tracked to determine when each loop closes.
56
+
57
+ Args:
58
+ ir_content: The IR content as a string.
59
+
60
+ Returns:
61
+ A list of tuples (start_line, end_line) for each scf.for loop found.
62
+ Line numbers are 0-indexed.
63
+ """
64
+ if not ir_content:
65
+ return []
66
+
67
+ loop_bounds: list[tuple[int, int]] = []
68
+ lines = ir_content.split("\n")
69
+
70
+ # Stack to track loop starts and their brace counts
71
+ # Each entry is (start_line, brace_count_at_start)
72
+ loop_stack: list[tuple[int, int]] = []
73
+ current_brace_count = 0
74
+
75
+ for line_idx, line in enumerate(lines):
76
+ # Check if this line starts a new scf.for loop
77
+ if "scf.for" in line:
78
+ loop_stack.append((line_idx, current_brace_count))
79
+
80
+ # Count braces on this line
81
+ for char in line:
82
+ if char == "{":
83
+ current_brace_count += 1
84
+ elif char == "}":
85
+ current_brace_count -= 1
86
+
87
+ # Check if we've closed any loops
88
+ while loop_stack and current_brace_count <= loop_stack[-1][1]:
89
+ start_line, _start_brace_count = loop_stack.pop()
90
+ # The loop ends at this line
91
+ loop_bounds.append((start_line, line_idx))
92
+
93
+ return loop_bounds
94
+
95
+
96
+ def find_inner_loop_bounds(ir_content: str) -> list[tuple[int, int]]:
97
+ """
98
+ Find the bounds of inner scf.for loops (loops without nested loops inside).
99
+
100
+ Inner loops are the primary candidates for Software Pipelining (SWP) as they
101
+ represent the innermost computation that can be optimized.
102
+
103
+ Args:
104
+ ir_content: The IR content as a string.
105
+
106
+ Returns:
107
+ A list of tuples (start_line, end_line) for each inner scf.for loop found.
108
+ Line numbers are 0-indexed.
109
+ """
110
+ all_loops = find_loop_bounds(ir_content)
111
+
112
+ if not all_loops:
113
+ return []
114
+
115
+ # Filter to keep only inner loops (loops that don't contain other loops)
116
+ inner_loops: list[tuple[int, int]] = []
117
+
118
+ for i, (start_i, end_i) in enumerate(all_loops):
119
+ # Check if any other loop is nested inside this loop
120
+ has_nested_loop = False
121
+ for j, (start_j, end_j) in enumerate(all_loops):
122
+ if i != j:
123
+ # Check if loop j is nested inside loop i
124
+ if start_i < start_j and end_j < end_i:
125
+ has_nested_loop = True
126
+ break
127
+
128
+ # If no nested loops found, this is an inner loop
129
+ if not has_nested_loop:
130
+ inner_loops.append((start_i, end_i))
131
+
132
+ return inner_loops
133
+
134
+
135
+ def find_loop_pipelining(
136
+ ttir_content: str,
137
+ ttgir_content: str,
138
+ ttir_loop_start: int,
139
+ ttir_loop_end: int,
140
+ loop_index: int,
141
+ ttir_to_ttgir_mapping: dict[str, dict],
142
+ ttgir_to_source_mapping: dict[str, dict],
143
+ python_source_content: str | None,
144
+ python_source_start_line: int,
145
+ ) -> dict[str, list[str]]:
146
+ """
147
+ Find pipelining information for a specific loop by identifying tt.load and tt.dot operations
148
+ in TTIR and mapping them to their corresponding operations in the original Python source code.
149
+
150
+ For each tt.load or tt.dot operation found in the TTIR loop, this function uses source
151
+ mappings to find the corresponding operations in TTGIR, then maps them back to the original
152
+ Python source code. Operations are categorized into three sections:
153
+ - prologue: Operations that appear before the loop body
154
+ - loop_body: Operations that appear within the loop body
155
+ - epilogue: Operations that appear after the loop body
156
+
157
+ Operations are merged together (both loads and dots) and sorted in program order
158
+ within each section.
159
+
160
+ Args:
161
+ ttir_content: The TTIR content as a string.
162
+ ttgir_content: The TTGIR content as a string.
163
+ ttir_loop_start: The starting line number of the loop in TTIR (0-indexed).
164
+ ttir_loop_end: The ending line number of the loop in TTIR (0-indexed).
165
+ ttir_to_ttgir_mapping: Source mapping from TTIR lines to TTGIR lines.
166
+ ttgir_to_source_mapping: Source mapping from TTGIR lines to original Python source.
167
+ python_source_content: The original Python source code content.
168
+
169
+ Returns:
170
+ A dictionary containing:
171
+ - "prologue": List of Python source line strings in program order
172
+ - "loop_body": List of Python source line strings in program order
173
+ - "epilogue": List of Python source line strings in program order
174
+ """
175
+ if not ttir_content or not ttgir_content:
176
+ return {
177
+ "prologue": [],
178
+ "loop_body": [],
179
+ "epilogue": [],
180
+ }
181
+
182
+ ttir_lines = ttir_content.split("\n")
183
+ ttgir_lines = ttgir_content.split("\n")
184
+ python_lines = python_source_content.split("\n") if python_source_content else []
185
+
186
+ def apply_trailing_space(op: str) -> str:
187
+ """
188
+ Add a trailing space to all ops to avoid false positives like
189
+ warp_group_dot and warp_group_dot_wait.
190
+ """
191
+ return op + " "
192
+
193
+ # Step 1: Find tt.load and tt.dot operations in TTIR loop
194
+ ttir_pipeline_lines: list[int] = []
195
+ pipeline_tt_ops = ["tt.load", "tt.dot"]
196
+ pipeline_tt_ops = [apply_trailing_space(op) for op in pipeline_tt_ops]
197
+ pipeline_ttgir_ops = [
198
+ "tt.load",
199
+ "tt.dot",
200
+ "async_copy_global_to_local",
201
+ "warp_group_dot",
202
+ ]
203
+ pipeline_ttgir_ops = [apply_trailing_space(op) for op in pipeline_ttgir_ops]
204
+ for line_idx in range(ttir_loop_start, min(ttir_loop_end + 1, len(ttir_lines))):
205
+ line = ttir_lines[line_idx]
206
+ for op in pipeline_tt_ops:
207
+ if op in line:
208
+ ttir_pipeline_lines.append(line_idx)
209
+ break
210
+
211
+ # Step 2: Find the corresponding loop in TTGIR using source mappings
212
+ # Map the TTIR loop bounds to TTGIR using source mappings
213
+ ttgir_inner_loops = find_inner_loop_bounds(ttgir_content)
214
+
215
+ if not ttgir_inner_loops:
216
+ # No loop found in TTGIR, return empty results
217
+ return {
218
+ "prologue": [],
219
+ "loop_body": [],
220
+ "epilogue": [],
221
+ }
222
+
223
+ # Use the first inner loop as the reference
224
+ # TODO: Implement more sophisticated mapping logic to match TTIR loops to TTGIR loops
225
+ ttgir_loop_start, ttgir_loop_end = ttgir_inner_loops[loop_index]
226
+
227
+ # Step 3: Map TTIR operations to TTGIR operations using source mappings
228
+ # and categorize them by their position relative to the TTGIR loop
229
+ # Store as (line_number, source_line) to maintain order before extracting just the source
230
+ prologue_ops: list[tuple[int, str]] = []
231
+ loop_body_ops: list[tuple[int, str]] = []
232
+ epilogue_ops: list[tuple[int, str]] = []
233
+
234
+ for ttir_line in ttir_pipeline_lines:
235
+ # Convert 0-indexed line to 1-indexed string key for mapping lookup
236
+ ttir_line_key = str(ttir_line + 1)
237
+
238
+ # Get the corresponding TTGIR lines from the source mapping
239
+ if ttir_line_key in ttir_to_ttgir_mapping:
240
+ ttgir_lines_list = ttir_to_ttgir_mapping[ttir_line_key].get(
241
+ "ttgir_lines", []
242
+ )
243
+
244
+ # For each mapped TTGIR line, categorize it
245
+ for ttgir_line in ttgir_lines_list:
246
+ # Convert back to 0-indexed
247
+ ttgir_line_idx = ttgir_line - 1
248
+
249
+ # Get the actual TTGIR line content to check if it's relevant
250
+ if ttgir_line_idx < len(ttgir_lines):
251
+ ttgir_source_line = ttgir_lines[ttgir_line_idx].strip()
252
+
253
+ # Only keep mappings to the "compute" op.
254
+ if any(op in ttgir_source_line for op in pipeline_ttgir_ops):
255
+ # Map TTGIR line back to Python source
256
+ ttgir_line_key = str(ttgir_line)
257
+ python_source_line = ttgir_source_line # Default to TTGIR line
258
+
259
+ if ttgir_line_key in ttgir_to_source_mapping:
260
+ source_info = ttgir_to_source_mapping[ttgir_line_key]
261
+ python_line_num = source_info.get("line")
262
+
263
+ if python_line_num and python_lines:
264
+ # Account for the offset: the Python source may not start at line 1
265
+ # python_line_num is the absolute line number in the original file
266
+ # python_source_start_line is where the extracted code starts
267
+ # So we need to subtract the offset to get the index in our python_lines array
268
+ python_line_idx = (
269
+ python_line_num - python_source_start_line
270
+ )
271
+ if 0 <= python_line_idx < len(python_lines):
272
+ python_source_line = python_lines[
273
+ python_line_idx
274
+ ].strip()
275
+
276
+ if ttgir_line_idx < ttgir_loop_start:
277
+ prologue_ops.append((ttgir_line_idx, python_source_line))
278
+ elif ttgir_loop_start <= ttgir_line_idx <= ttgir_loop_end:
279
+ loop_body_ops.append((ttgir_line_idx, python_source_line))
280
+ else:
281
+ epilogue_ops.append((ttgir_line_idx, python_source_line))
282
+
283
+ # Step 4: Sort each section by line number to maintain program order
284
+ prologue_ops.sort(key=lambda x: x[0])
285
+ loop_body_ops.sort(key=lambda x: x[0])
286
+ epilogue_ops.sort(key=lambda x: x[0])
287
+
288
+ # Extract just the source lines (without line numbers)
289
+ prologue_lines = [line for _, line in prologue_ops]
290
+ loop_body_lines = [line for _, line in loop_body_ops]
291
+ epilogue_lines = [line for _, line in epilogue_ops]
292
+
293
+ # Log the pipelining results
294
+ logger.info(
295
+ f"Loop pipelining results (TTIR lines {ttir_loop_start}-{ttir_loop_end}):"
296
+ )
297
+ logger.info(f" Prologue ({len(prologue_lines)} ops):")
298
+ for line in prologue_lines:
299
+ logger.info(f" {line}")
300
+ logger.info(f" Loop Body ({len(loop_body_lines)} ops):")
301
+ for line in loop_body_lines:
302
+ logger.info(f" {line}")
303
+ logger.info(f" Epilogue ({len(epilogue_lines)} ops):")
304
+ for line in epilogue_lines:
305
+ logger.info(f" {line}")
306
+
307
+ return {
308
+ "prologue": prologue_lines,
309
+ "loop_body": loop_body_lines,
310
+ "epilogue": epilogue_lines,
311
+ }
312
+
313
+
314
+ def generate_loop_schedule(
315
+ ttir_key: str,
316
+ ttgir_key: str,
317
+ file_content: dict[str, str],
318
+ file_path: dict[str, str],
319
+ source_mappings: dict[str, dict],
320
+ python_source_content: str | None,
321
+ python_source_start_line: int,
322
+ ) -> list[dict]:
323
+ """
324
+ Generate loop schedule information by finding inner scf.for loops in TTIR
325
+ and analyzing their pipelining potential using source mappings.
326
+
327
+ Only inner loops (loops without nested loops) are considered as they are
328
+ the primary candidates for Software Pipelining (SWP).
329
+
330
+ Args:
331
+ ttir_key: Key for the TTIR file.
332
+ ttgir_key: Key for the TTGIR file.
333
+ file_content: Dictionary mapping file keys to content.
334
+ file_path: Dictionary mapping file keys to file paths.
335
+ source_mappings: Dictionary containing source mappings between IR stages.
336
+ python_source_content: The original Python source code content.
337
+ python_source_start_line: The starting line number of the Python source in the original file.
338
+
339
+ Returns:
340
+ A list of dictionaries, each containing:
341
+ - "loop_bounds": Tuple of (start_line, end_line) for the loop in TTIR
342
+ - "pipelining": Dictionary with Python source lines for operations
343
+ """
344
+ ttir_content = load_ir_contents(ttir_key, file_content, file_path)
345
+ ttgir_content = load_ir_contents(ttgir_key, file_content, file_path)
346
+
347
+ # Get the TTIR to TTGIR mapping and TTGIR to source mapping
348
+ ttir_to_ttgir_mapping = source_mappings.get("ttir", {})
349
+ ttgir_to_source_mapping = source_mappings.get("ttgir", {})
350
+
351
+ # Find only inner loops (loops without nested loops inside)
352
+ inner_loop_bounds = find_inner_loop_bounds(ttir_content)
353
+ # TODO: Fix loop mapping with multiple loops.
354
+ inner_loop_bounds = inner_loop_bounds[:1]
355
+
356
+ # For each inner loop, find pipelining information
357
+ loop_schedules = []
358
+ for i, (loop_start, loop_end) in enumerate(inner_loop_bounds):
359
+ pipelining_info = find_loop_pipelining(
360
+ ttir_content,
361
+ ttgir_content,
362
+ loop_start,
363
+ loop_end,
364
+ i,
365
+ ttir_to_ttgir_mapping,
366
+ ttgir_to_source_mapping,
367
+ python_source_content,
368
+ python_source_start_line,
369
+ )
370
+ loop_schedules.append(pipelining_info)
371
+
372
+ return loop_schedules
373
+
374
+
49
375
  def _generate_ir_analysis(entry: str):
50
376
  payload = entry.setdefault("payload", {})
51
377
  file_content = payload.get("file_content", {})
52
378
  file_path = payload.get("file_path", {})
379
+ source_mappings = payload.get("source_mappings", {})
53
380
 
54
381
  # Find the IR file keys
382
+ ttir_key = next((k for k in file_content if k.endswith(".ttir")), None)
55
383
  ttgir_key = next((k for k in file_content if k.endswith(".ttgir")), None)
56
384
  amdgcn_key = next((k for k in file_content if k.endswith(".amdgcn")), None)
57
385
  # Skip if no IR files found
58
- if not (ttgir_key or amdgcn_key):
59
- logger.debug("No AMD IR found")
386
+ if not (ttir_key or ttgir_key or amdgcn_key):
387
+ logger.debug("No IR found")
60
388
  return {}
61
389
  ir_analysis = {}
62
- if amdgcn_key:
390
+ if amdgcn_key and ttgir_key:
391
+ # Add BufferOps information
63
392
  ttgir_bufferops_info = process_amd_ttgir_bufferops(
64
393
  ttgir_key, file_content, file_path
65
394
  )
@@ -74,4 +403,25 @@ def _generate_ir_analysis(entry: str):
74
403
  io_counts["amd_gcn_bufferops_count"] = gcn_bufferops_info
75
404
  if io_counts:
76
405
  ir_analysis["io_counts"] = io_counts
406
+ if ttir_key and ttgir_key:
407
+ # Get Python source content and start line if available
408
+ python_source_content = None
409
+ python_source_start_line = 1 # Default to 1 if not available
410
+ python_source_info = payload.get("python_source")
411
+ if python_source_info:
412
+ python_source_content = python_source_info.get("code")
413
+ python_source_start_line = python_source_info.get("start_line", 1)
414
+
415
+ # Add loop schedule information
416
+ loop_schedule = generate_loop_schedule(
417
+ ttir_key,
418
+ ttgir_key,
419
+ file_content,
420
+ file_path,
421
+ source_mappings,
422
+ python_source_content,
423
+ python_source_start_line,
424
+ )
425
+ if loop_schedule:
426
+ ir_analysis["loop_schedules"] = loop_schedule
77
427
  return ir_analysis
tritonparse/ir_parser.py CHANGED
@@ -44,6 +44,14 @@ ALIAS_WITH_NAME_PATTERN = re.compile(
44
44
  # Example: #loc20 = loc(#loc16)
45
45
  ALIAS_SIMPLE_PATTERN = re.compile(r"#loc(\d+)\s*=\s*loc\(\s*#loc(\d*)\s*\)")
46
46
 
47
+ # Callsite loc definitions in TTIR/TTGIR
48
+ # Example: #loc220 = loc(callsite(#loc57 at #loc190))
49
+ # Captures: loc_id, callee_loc_id, caller_loc_id
50
+ # Note: Uses (\d*) to match optional numbers (for bare #loc references)
51
+ CALLSITE_PATTERN = re.compile(
52
+ r"#loc(\d+)\s*=\s*loc\(\s*callsite\(\s*#loc(\d*)\s+at\s+#loc(\d*)\s*\)\s*\)"
53
+ )
54
+
47
55
 
48
56
  def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
49
57
  """
@@ -141,6 +149,50 @@ def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
141
149
  for alias_id, target_id in alias_map.items():
142
150
  if alias_id not in locations:
143
151
  resolve_alias(alias_id)
152
+
153
+ # Collect callsite definitions
154
+ callsite_defs = []
155
+ for i, line in enumerate(ir_content.split("\n"), start=1):
156
+ if m := CALLSITE_PATTERN.search(line):
157
+ loc_id, callee_id, caller_id = m.groups()
158
+ # Empty strings map to main loc key ""
159
+ callsite_defs.append((loc_id, callee_id or "", caller_id or "", i))
160
+
161
+ # Resolve callsite definitions
162
+ # A callsite inherits the location from its callee (the code being called)
163
+ # and stores a reference to its caller (the code doing the calling)
164
+ for loc_id, callee_id, caller_id, def_line in callsite_defs:
165
+ if loc_id not in locations: # Avoid overwriting existing definitions
166
+ if callee_id in locations:
167
+ # Inherit location info from callee
168
+ callee_info = locations[callee_id]
169
+ locations[loc_id] = {
170
+ "file": callee_info["file"],
171
+ "line": callee_info["line"],
172
+ "column": callee_info["column"],
173
+ "def_line": def_line,
174
+ "is_callsite": True,
175
+ "callsite_callee": callee_id,
176
+ "callsite_caller": caller_id,
177
+ }
178
+ else:
179
+ logger.warning(
180
+ f"Callsite #loc{loc_id} references undefined callee #loc{callee_id}"
181
+ )
182
+ # Note: We don't add this callsite to locations since callee is missing
183
+
184
+ # Verify caller references (warning only, don't block)
185
+ for loc_id, _callee_id, caller_id, _def_line in callsite_defs:
186
+ if loc_id in locations and caller_id and caller_id not in locations:
187
+ logger.warning(
188
+ f"Callsite #loc{loc_id} references undefined caller #loc{caller_id}"
189
+ )
190
+
191
+ # Attach definition line and alias metadata
192
+ for k, v in def_line_map.items():
193
+ if k in locations:
194
+ locations[k]["def_line"] = v
195
+ for alias_id, target_id in alias_map.items():
144
196
  if alias_id in locations:
145
197
  locations[alias_id]["alias_of"] = target_id
146
198
  if alias_id in alias_name_map:
@@ -48,6 +48,14 @@ TRITON_TRACE_LAUNCH = os.getenv("TRITON_TRACE_LAUNCH", None) in ["1", "true", "T
48
48
  TRITONPARSE_MORE_TENSOR_INFORMATION = os.getenv(
49
49
  "TRITONPARSE_MORE_TENSOR_INFORMATION", None
50
50
  ) in ["1", "true", "True"]
51
+ # Enable full Python source file extraction instead of just the function definition
52
+ TRITON_FULL_PYTHON_SOURCE = os.getenv("TRITON_FULL_PYTHON_SOURCE", "0") in [
53
+ "1",
54
+ "true",
55
+ "True",
56
+ ]
57
+ # Maximum file size for full source extraction (default 10MB)
58
+ TRITON_MAX_SOURCE_SIZE = int(os.getenv("TRITON_MAX_SOURCE_SIZE", str(10 * 1024 * 1024)))
51
59
  # Inductor compiled kernel's launch tracing needs this flag to be set.
52
60
  # If TRITON_TRACE_LAUNCH is enabled, also enable TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK
53
61
  TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK = (
@@ -727,6 +735,17 @@ def extract_python_source_info(trace_data: Dict[str, Any], source):
727
735
  from the provided source object (typically an ASTSource or IRSource instance).
728
736
  It adds file path, line numbers, and the actual source code to the trace_data.
729
737
 
738
+ By default, only the function definition is extracted. Set TRITON_FULL_PYTHON_SOURCE=1
739
+ to extract the entire Python source file.
740
+ @TODO: we should enable it by default in next diff and track the compilation time regression
741
+
742
+ Environment Variables:
743
+ TRITON_FULL_PYTHON_SOURCE: If set to "1", extract the full Python file
744
+ instead of just the function definition.
745
+ TRITON_MAX_SOURCE_SIZE: Maximum file size in bytes for full source extraction
746
+ (default: 10MB). Files larger than this will fall back
747
+ to function-only mode.
748
+
730
749
  Args:
731
750
  trace_data (Dict[str, Any]): Dictionary to store extracted information
732
751
  source (Union[ASTSource, IRSource]): Source object containing kernel function information
@@ -738,23 +757,77 @@ def extract_python_source_info(trace_data: Dict[str, Any], source):
738
757
  if isinstance(source, IRSource):
739
758
  return
740
759
 
741
- # Get the original Python source code for the kernel
760
+ # Get the function reference
761
+ if isinstance(fn := source.fn, JITFunction):
762
+ fn_ref = fn.fn
763
+ else:
764
+ fn_ref = source.fn
765
+
766
+ python_source_file = inspect.getfile(fn_ref)
767
+
768
+ # Get function range information
742
769
  if (
743
770
  isinstance(fn := source.fn, JITFunction)
744
771
  and hasattr(fn, "starting_line_number")
745
772
  and hasattr(fn, "raw_src")
746
773
  ):
747
- start_line_number = fn.starting_line_number
774
+ function_start_line = fn.starting_line_number
748
775
  source_lines = fn.raw_src
749
776
  else:
750
- source_lines, start_line_number = inspect.getsourcelines(fn.fn)
777
+ source_lines, function_start_line = inspect.getsourcelines(fn_ref)
778
+
779
+ function_end_line = function_start_line + len(source_lines) - 1
780
+
781
+ if TRITON_FULL_PYTHON_SOURCE:
782
+ # Full file mode: read the entire Python file
783
+ try:
784
+ # Check file size before reading
785
+ file_size = os.path.getsize(python_source_file)
786
+ except OSError as e:
787
+ log.warning(
788
+ f"Failed to check file size for {python_source_file}: {e}. "
789
+ f"Falling back to function-only mode."
790
+ )
791
+ use_full_source = False
792
+ else:
793
+ if file_size > TRITON_MAX_SOURCE_SIZE:
794
+ log.warning(
795
+ f"Source file {python_source_file} is too large ({file_size} bytes, "
796
+ f"limit: {TRITON_MAX_SOURCE_SIZE} bytes). Falling back to function-only mode."
797
+ )
798
+ use_full_source = False
799
+ else:
800
+ use_full_source = True
801
+
802
+ if use_full_source:
803
+ try:
804
+ with open(python_source_file, "r", encoding="utf-8") as f:
805
+ file_content = f.read()
806
+
807
+ # Calculate total lines
808
+ total_lines = len(file_content.split("\n"))
809
+
810
+ trace_data["python_source"] = {
811
+ "file_path": python_source_file,
812
+ "start_line": 1,
813
+ "end_line": total_lines,
814
+ "code": file_content,
815
+ # Add function range for frontend highlighting and scrolling
816
+ "function_start_line": function_start_line,
817
+ "function_end_line": function_end_line,
818
+ }
819
+ return
820
+ except (OSError, UnicodeDecodeError) as e:
821
+ log.warning(
822
+ f"Failed to read full source file {python_source_file}: {e}. "
823
+ f"Falling back to function-only mode."
824
+ )
751
825
 
752
- python_source_file = inspect.getfile(fn.fn)
753
- end_line_number = start_line_number + len(source_lines)
826
+ # Default behavior: only extract function definition
754
827
  trace_data["python_source"] = {
755
828
  "file_path": python_source_file,
756
- "start_line": start_line_number,
757
- "end_line": end_line_number,
829
+ "start_line": function_start_line,
830
+ "end_line": function_end_line,
758
831
  "code": "".join(source_lines),
759
832
  }
760
833
 
@@ -910,7 +983,7 @@ class TritonTraceHandler(logging.StreamHandler):
910
983
  )
911
984
  elif not os.access(TRACE_LOG_DIR, os.W_OK):
912
985
  log.info(
913
- "TritonTraceHandler: disabled because %s is not writeable",
986
+ "TritonTraceHandler: disabled because %s is not writable",
914
987
  TRACE_LOG_DIR,
915
988
  )
916
989
  else:
@@ -77,6 +77,11 @@ def generate_source_mappings(
77
77
  "column": info["column"],
78
78
  f"{ir_type}_line": ln,
79
79
  }
80
+ # Propagate callsite metadata if present
81
+ if info.get("is_callsite"):
82
+ entry["is_callsite"] = True
83
+ entry["callsite_callee"] = info["callsite_callee"]
84
+ entry["callsite_caller"] = info["callsite_caller"]
80
85
  # Propagate alias metadata if present
81
86
  if "alias_name" in info:
82
87
  entry["alias_name"] = info["alias_name"]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tritonparse
3
- Version: 0.3.1.dev20251029071541
3
+ Version: 0.3.1.dev20251031071507
4
4
  Summary: TritonParse: A Compiler Tracer, Visualizer, and mini-Reproducer Generator for Triton Kernels
5
5
  Author-email: Yueming Hao <yhao@meta.com>
6
6
  License-Expression: BSD-3-Clause
@@ -5,15 +5,15 @@ tritonparse/common.py,sha256=MJo9bVCgSKkwXpEoUkUczPo_5jOYpJgXLq4UsWYqN3c,13924
5
5
  tritonparse/context_manager.py,sha256=OdMn11qbApYL2c9IlbUpcT27r04ZSa4DfvrY2mLA958,2243
6
6
  tritonparse/event_diff.py,sha256=USCjfjYr-7Ie-EfZgtCFMZMA1KRzFRDe7yDFy98zYI4,4962
7
7
  tritonparse/extract_source_mappings.py,sha256=Z6UxFj2cCE5NCWLQTYPKqUpLfbYhqP8xgCl5mvud9KI,1451
8
- tritonparse/ir_analysis.py,sha256=MoOXuHsUGZ705R4JnXmlsrBn9gJdLO1Dnf0L5AxcaBM,2551
9
- tritonparse/ir_parser.py,sha256=MH4RwoNZMBdWUxkFyEhemJ7Aa7-asoba66b06bGPNsk,13237
8
+ tritonparse/ir_analysis.py,sha256=DZz9H8DqW753UkYECnyt6ATC6J1yvLxOHVRHHAZbrVg,16627
9
+ tritonparse/ir_parser.py,sha256=JQ7hsevmhFGmtZ3CoXi4utcomAycBQTT-KFjSva2K8U,15565
10
10
  tritonparse/mapper.py,sha256=QBCUMHM9pu3x3ahFp0wyXJbmv9TFGVPdkcLULok1E-k,4205
11
11
  tritonparse/shared_vars.py,sha256=RifXq55KisHgspYAmGcaCWY6ZHX8iejFHvwIewvcWZE,707
12
12
  tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
13
13
  tritonparse/sourcemap_utils.py,sha256=uI02n5Sgnlx7Nc15QAX5N6_tZZMips0PyJuo1n3eouY,2654
14
- tritonparse/structured_logging.py,sha256=L1xkkCx8Jr9YQbM0Kgtf2g6L3aWMkYOEeFFEOSo8Lkk,60306
14
+ tritonparse/structured_logging.py,sha256=aBTfRVuZzd7YJV-Fgg513SPlbEceyWosONLqaTZS69k,63319
15
15
  tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
16
- tritonparse/trace_processor.py,sha256=aQPqlnpTtWoGzHYv4BXWUH4nCeUQGSK3o-fj0LD9I0c,14147
16
+ tritonparse/trace_processor.py,sha256=AW4YDrPDayURtmePkFi5m5p6P7OTi1UlTPbbrPzujwY,14418
17
17
  tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
18
18
  tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
19
19
  tritonparse/reproducer/cli.py,sha256=95AgH9QOlSFpkC5iR89XV9wubv_5vfD1MKl2IxpAIzs,1718
@@ -34,9 +34,9 @@ tritonparse/tools/format_fix.py,sha256=ISalg_N_L7Xktag3mLr-G9T6Opxv793s1WG6A9wUt
34
34
  tritonparse/tools/load_tensor.py,sha256=7-LbpboKDNJFBLNhiKS3enoqRlVAb55OjPc70PwHXAw,2789
35
35
  tritonparse/tools/prettify_ndjson.py,sha256=kR8hmBCv-iJeuzpi2_6CZv9T4_edRQbBOSOPpMm6wrw,11117
36
36
  tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
37
- tritonparse-0.3.1.dev20251029071541.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
38
- tritonparse-0.3.1.dev20251029071541.dist-info/METADATA,sha256=IX0-TFnQLPgqNePCfPvAn87E9HbfQWXH7mns7P2UBp0,8282
39
- tritonparse-0.3.1.dev20251029071541.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
40
- tritonparse-0.3.1.dev20251029071541.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
41
- tritonparse-0.3.1.dev20251029071541.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
42
- tritonparse-0.3.1.dev20251029071541.dist-info/RECORD,,
37
+ tritonparse-0.3.1.dev20251031071507.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
38
+ tritonparse-0.3.1.dev20251031071507.dist-info/METADATA,sha256=NA-UNdnVGTk0HMC3iixsPsHYYK_7TaBCiEdTI3PYBaU,8282
39
+ tritonparse-0.3.1.dev20251031071507.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
40
+ tritonparse-0.3.1.dev20251031071507.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
41
+ tritonparse-0.3.1.dev20251031071507.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
42
+ tritonparse-0.3.1.dev20251031071507.dist-info/RECORD,,