tritonparse 0.3.1.dev20251028071524__tar.gz → 0.3.1.dev20251029071541__tar.gz

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.

Files changed (119) hide show
  1. {tritonparse-0.3.1.dev20251028071524/tritonparse.egg-info → tritonparse-0.3.1.dev20251029071541}/PKG-INFO +1 -1
  2. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/cli.py +8 -0
  3. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/cli.py +5 -0
  4. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/orchestrator.py +2 -2
  5. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/placeholder_replacer.py +35 -17
  6. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/templates/example.py +5 -1
  7. tritonparse-0.3.1.dev20251029071541/tritonparse/reproducer/templates/tritonbench.py +103 -0
  8. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/utils.py +7 -2
  9. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541/tritonparse.egg-info}/PKG-INFO +1 -1
  10. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse.egg-info/SOURCES.txt +2 -0
  11. tritonparse-0.3.1.dev20251029071541/website/src/components/CodeViewer.css +43 -0
  12. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/CodeViewer.tsx +65 -3
  13. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/utils/dataLoader.ts +6 -0
  14. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.ci/README.md +0 -0
  15. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.ci/install-project.sh +0 -0
  16. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.ci/install-triton.sh +0 -0
  17. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.ci/run-tests.sh +0 -0
  18. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.ci/setup.sh +0 -0
  19. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.github/PAGES_SETUP.md +0 -0
  20. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.github/workflows/deploy-pages-standalone.yml +0 -0
  21. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.github/workflows/deploy-pages.yml +0 -0
  22. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.github/workflows/nightly-pypi.yml +0 -0
  23. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.github/workflows/test.yml +0 -0
  24. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/.gitignore +0 -0
  25. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/CHANGELOG.md +0 -0
  26. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/CODE_OF_CONDUCT.md +0 -0
  27. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/CONTRIBUTING.md +0 -0
  28. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/LICENSE +0 -0
  29. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/Makefile +0 -0
  30. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/README.md +0 -0
  31. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/__init__.py +0 -0
  32. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/pyproject.toml +0 -0
  33. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/run.py +0 -0
  34. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/setup.cfg +0 -0
  35. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/README.md +0 -0
  36. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/__init__.py +0 -0
  37. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/logs/dedicated_log_triton_trace_findhao_.ndjson +0 -0
  38. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/parsed_output/dedicated_log_triton_trace_findhao__mapped.ndjson.gz +0 -0
  39. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/parsed_output/f0_fc0_a0_cai-.ndjson.gz +0 -0
  40. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/parsed_output/log_file_list.json +0 -0
  41. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz +0 -0
  42. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/example_output/parsed_output_complex/log_file_list.json +0 -0
  43. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/test_add.py +0 -0
  44. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tests/test_tritonparse.py +0 -0
  45. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/__init__.py +0 -0
  46. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/__main__.py +0 -0
  47. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/common.py +0 -0
  48. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/context_manager.py +0 -0
  49. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/event_diff.py +0 -0
  50. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/extract_source_mappings.py +0 -0
  51. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/ir_analysis.py +0 -0
  52. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/ir_parser.py +0 -0
  53. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/mapper.py +0 -0
  54. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/__init__.py +0 -0
  55. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/function_extractor.py +0 -0
  56. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/ingestion/ndjson.py +0 -0
  57. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/templates/__init__.py +0 -0
  58. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/templates/loader.py +0 -0
  59. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/types.py +0 -0
  60. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/shared_vars.py +0 -0
  61. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/source_type.py +0 -0
  62. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/sourcemap_utils.py +0 -0
  63. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/structured_logging.py +0 -0
  64. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/__init__.py +0 -0
  65. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/decompress_bin_ndjson.py +0 -0
  66. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/disasm.py +0 -0
  67. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/format_fix.py +0 -0
  68. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/load_tensor.py +0 -0
  69. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/prettify_ndjson.py +0 -0
  70. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tools/readme.md +0 -0
  71. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/tp_logger.py +0 -0
  72. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/trace_processor.py +0 -0
  73. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse/utils.py +0 -0
  74. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse.egg-info/dependency_links.txt +0 -0
  75. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse.egg-info/entry_points.txt +0 -0
  76. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse.egg-info/requires.txt +0 -0
  77. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/tritonparse.egg-info/top_level.txt +0 -0
  78. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/eslint.config.js +0 -0
  79. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/index.html +0 -0
  80. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/package-lock.json +0 -0
  81. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/package.json +0 -0
  82. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/public/dedicated_log_triton_trace_findhao__mapped.ndjson.gz +0 -0
  83. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/public/f0_fc0_a0_cai-.ndjson +0 -0
  84. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/public/favicon.ico +0 -0
  85. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/public/logo.svg +0 -0
  86. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/scripts/inline-html.js +0 -0
  87. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/scripts/update_deps.sh +0 -0
  88. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/App.css +0 -0
  89. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/App.tsx +0 -0
  90. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/assets/react.svg +0 -0
  91. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/ArgumentViewer.tsx +0 -0
  92. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/Callstack.tsx +0 -0
  93. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/CodeComparisonView.tsx +0 -0
  94. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/CompilationInfo.tsx +0 -0
  95. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/CopyCodeButton.tsx +0 -0
  96. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/DataSourceSelector.tsx +0 -0
  97. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/DiffComparisonView.tsx +0 -0
  98. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/DiffViewer.tsx +0 -0
  99. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/ExternalLink.tsx +0 -0
  100. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/SingleCodeViewer.tsx +0 -0
  101. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/StackDiffViewer.tsx +0 -0
  102. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/ToggleSwitch.tsx +0 -0
  103. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/TritonIRs.tsx +0 -0
  104. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/components/WelcomeScreen.tsx +0 -0
  105. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/context/FileDiffSession.tsx +0 -0
  106. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/index.css +0 -0
  107. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/main.tsx +0 -0
  108. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/pages/CodeView.tsx +0 -0
  109. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/pages/FileDiffView.tsx +0 -0
  110. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/pages/IRAnalysis.tsx +4 -4
  111. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/pages/KernelOverview.tsx +0 -0
  112. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/utils/fbDetection.ts +0 -0
  113. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/utils/safeImport.ts +0 -0
  114. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/utils/tensor.ts +0 -0
  115. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/src/vite-env.d.ts +0 -0
  116. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/tsconfig.app.json +0 -0
  117. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/tsconfig.json +0 -0
  118. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/tsconfig.node.json +0 -0
  119. {tritonparse-0.3.1.dev20251028071524 → tritonparse-0.3.1.dev20251029071541}/website/vite.config.ts +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tritonparse
3
- Version: 0.3.1.dev20251028071524
3
+ Version: 0.3.1.dev20251029071541
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
@@ -68,12 +68,20 @@ def main():
68
68
  }
69
69
  unified_parse(**parse_args)
70
70
  elif args.func == "reproduce":
71
+ replacer = None
72
+ if args.use_fbcode:
73
+ from tritonparse.fb.reproducer.replacer import FBCodePlaceholderReplacer
74
+
75
+ replacer = FBCodePlaceholderReplacer()
76
+ print(f"Using FBCode placeholder replacer for template: {args.template}")
77
+
71
78
  reproduce(
72
79
  input_path=args.input,
73
80
  line_index=args.line - 1, # Convert 1-based line number to 0-based index
74
81
  out_dir=args.out_dir,
75
82
  template=args.template,
76
83
  kernel_import=args.kernel_import,
84
+ replacer=replacer,
77
85
  )
78
86
  else:
79
87
  raise RuntimeError(f"Unknown command: {args.func}")
@@ -46,3 +46,8 @@ def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
46
46
  "Defaults to 'default'."
47
47
  ),
48
48
  )
49
+ parser.add_argument(
50
+ "--use-fbcode",
51
+ action="store_true",
52
+ help=("Use fbcode to setup repro environment."),
53
+ )
@@ -23,7 +23,7 @@ def reproduce(
23
23
  template: str,
24
24
  replacer: Optional[PlaceholderReplacer] = None,
25
25
  kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
26
- ) -> dict[str, Path]:
26
+ ) -> dict[str, str]:
27
27
  """
28
28
  Generate a reproducer script from NDJSON trace file.
29
29
 
@@ -45,7 +45,7 @@ def reproduce(
45
45
  f"Built context bundle for kernel: {context_bundle.kernel_info.function_name}"
46
46
  )
47
47
  out_py_path, temp_json_path = determine_output_paths(
48
- out_dir, context_bundle.kernel_info.function_name
48
+ out_dir, context_bundle.kernel_info.function_name, template
49
49
  )
50
50
  save_prettified_json(context_bundle.raw_launch_event, temp_json_path)
51
51
 
@@ -76,21 +76,39 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
76
76
  - # {{KERNEL_INVOCATION_PLACEHOLDER}}: Replaced with kernel invocation code
77
77
  """
78
78
 
79
+ KERNEL_NAME_PLACEHOLDER = "{{KERNEL_NAME_PLACEHOLDER}}"
80
+ JSON_FILE_NAME_PLACEHOLDER = "{{JSON_FILE_NAME_PLACEHOLDER}}"
81
+ IR_OVERRIDE_SETUP_PLACEHOLDER = "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}"
82
+ KERNEL_SYSPATH_PLACEHOLDER = "# {{KERNEL_SYSPATH_PLACEHOLDER}}"
83
+ KERNEL_IMPORT_PLACEHOLDER = "# {{KERNEL_IMPORT_PLACEHOLDER}}"
84
+ UTILITY_FUNCTIONS_PLACEHOLDER = "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}"
85
+ KERNEL_INVOCATION_PLACEHOLDER = "# {{KERNEL_INVOCATION_PLACEHOLDER}}"
86
+
79
87
  def __init__(self):
80
88
  super().__init__()
81
89
  # Register all default handlers
82
- self.register("{{JSON_FILE_NAME_PLACEHOLDER}}", self._replace_json_filename)
90
+ self.register(self.JSON_FILE_NAME_PLACEHOLDER, self._replace_json_filename)
83
91
  self.register(
84
- "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", self._replace_ir_override_setup
92
+ self.IR_OVERRIDE_SETUP_PLACEHOLDER, self._replace_ir_override_setup
85
93
  )
86
- self.register("# {{KERNEL_SYSPATH_PLACEHOLDER}}", self._replace_kernel_syspath)
87
- self.register("# {{KERNEL_IMPORT_PLACEHOLDER}}", self._replace_kernel_import)
94
+ self.register(self.KERNEL_SYSPATH_PLACEHOLDER, self._replace_kernel_syspath)
95
+ self.register(self.KERNEL_IMPORT_PLACEHOLDER, self._replace_kernel_import)
88
96
  self.register(
89
- "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", self._replace_utility_functions
97
+ self.UTILITY_FUNCTIONS_PLACEHOLDER, self._replace_utility_functions
90
98
  )
91
99
  self.register(
92
- "# {{KERNEL_INVOCATION_PLACEHOLDER}}", self._replace_kernel_invocation
100
+ self.KERNEL_INVOCATION_PLACEHOLDER, self._replace_kernel_invocation
93
101
  )
102
+ self.register(self.KERNEL_NAME_PLACEHOLDER, self._replace_kernel_name)
103
+
104
+ def _replace_kernel_name(
105
+ self, code: str, context_bundle: ContextBundle, **kwargs
106
+ ) -> str:
107
+ """Replace the kernel name placeholder."""
108
+ kernel_name = context_bundle.kernel_info.function_name
109
+ if not kernel_name:
110
+ raise ValueError("Kernel function name is not available")
111
+ return code.replace(self.KERNEL_NAME_PLACEHOLDER, kernel_name)
94
112
 
95
113
  def _replace_json_filename(
96
114
  self, code: str, context_bundle: ContextBundle, **kwargs
@@ -99,7 +117,7 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
99
117
  temp_json_path = kwargs.get("temp_json_path")
100
118
  if temp_json_path is None:
101
119
  raise ValueError("temp_json_path is required for JSON filename replacement")
102
- return code.replace("{{JSON_FILE_NAME_PLACEHOLDER}}", temp_json_path.name)
120
+ return code.replace(self.JSON_FILE_NAME_PLACEHOLDER, temp_json_path.name)
103
121
 
104
122
  def _replace_ir_override_setup(
105
123
  self, code: str, context_bundle: ContextBundle, **kwargs
@@ -108,7 +126,7 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
108
126
  kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
109
127
 
110
128
  if kernel_import != KernelImportMode.OVERRIDE_TTIR:
111
- return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", "")
129
+ return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, "")
112
130
 
113
131
  comp_json_filename = kwargs.get("comp_json_filename")
114
132
  if not comp_json_filename:
@@ -158,7 +176,7 @@ _original_autotune = triton.autotune
158
176
  triton.autotune = _patched_autotune
159
177
  '''
160
178
 
161
- return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", setup_code)
179
+ return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, setup_code)
162
180
 
163
181
  def _replace_kernel_syspath(
164
182
  self, code: str, context_bundle: ContextBundle, **kwargs
@@ -168,15 +186,15 @@ triton.autotune = _patched_autotune
168
186
 
169
187
  if kernel_import == KernelImportMode.DEFAULT:
170
188
  sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
171
- return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", sys_stmt)
189
+ return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, sys_stmt)
172
190
  elif kernel_import == KernelImportMode.COPY:
173
191
  comment = (
174
192
  "# Kernel sys.path setup skipped - kernel source code embedded below"
175
193
  )
176
- return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", comment)
194
+ return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
177
195
  elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
178
196
  comment = "# Kernel sys.path setup skipped - using IR override mode"
179
- return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", comment)
197
+ return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
180
198
  else:
181
199
  raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
182
200
 
@@ -190,7 +208,7 @@ triton.autotune = _patched_autotune
190
208
  _, import_statement = _generate_import_statements(
191
209
  context_bundle.kernel_info
192
210
  )
193
- return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", import_statement)
211
+ return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, import_statement)
194
212
  elif kernel_import == KernelImportMode.COPY:
195
213
  source_code = context_bundle.kernel_info.source_code
196
214
  func_name = context_bundle.kernel_info.function_name
@@ -216,10 +234,10 @@ triton.autotune = _patched_autotune
216
234
  embedded_code += "\n" + source_code
217
235
  embedded_code += f"\n\n# Use kernel function directly\nimported_kernel_function = {func_name}"
218
236
 
219
- return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", embedded_code)
237
+ return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
220
238
  elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
221
239
  comment = "# Kernel import skipped - using IR override mode with TTIR"
222
- return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", comment)
240
+ return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, comment)
223
241
  else:
224
242
  raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
225
243
 
@@ -228,7 +246,7 @@ triton.autotune = _patched_autotune
228
246
  ) -> str:
229
247
  """Replace the utility functions placeholder with extracted functions."""
230
248
  utility_code = extract_utility_functions()
231
- return code.replace("# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", utility_code)
249
+ return code.replace(self.UTILITY_FUNCTIONS_PLACEHOLDER, utility_code)
232
250
 
233
251
  def _replace_kernel_invocation(
234
252
  self, code: str, context_bundle: ContextBundle, **kwargs
@@ -237,4 +255,4 @@ triton.autotune = _patched_autotune
237
255
  source_code = context_bundle.kernel_info.source_code
238
256
  pos_args, kw_args = _parse_kernel_signature(source_code)
239
257
  invocation_snippet = _generate_invocation_snippet(pos_args, kw_args)
240
- return code.replace("# {{KERNEL_INVOCATION_PLACEHOLDER}}", invocation_snippet)
258
+ return code.replace(self.KERNEL_INVOCATION_PLACEHOLDER, invocation_snippet)
@@ -14,7 +14,7 @@ import torch
14
14
  # {{UTILITY_FUNCTIONS_PLACEHOLDER}}
15
15
 
16
16
 
17
- if __name__ == "__main__":
17
+ def launch_kernel():
18
18
  script_dir = Path(__file__).resolve().parent # noqa: F821
19
19
  json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
20
20
  grid, args_dict = create_args_from_json_file(str(json_file)) # noqa: F821
@@ -28,3 +28,7 @@ if __name__ == "__main__":
28
28
 
29
29
  torch.cuda.synchronize()
30
30
  print("Kernel execution finished.")
31
+
32
+
33
+ if __name__ == "__main__":
34
+ launch_kernel()
@@ -0,0 +1,103 @@
1
+ # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
2
+
3
+ from pathlib import Path
4
+ from typing import Any, Callable, Dict, Optional, Tuple
5
+
6
+ import torch
7
+ from tritonbench.utils.triton_op import (
8
+ BenchmarkOperator,
9
+ register_benchmark,
10
+ REGISTERED_X_VALS,
11
+ )
12
+
13
+
14
+ imported_kernel_function: Optional[Callable[[Tuple[int], Dict[str, Any]], None]] = None
15
+
16
+ # {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
17
+
18
+ # {{KERNEL_SYSPATH_PLACEHOLDER}}
19
+
20
+ # {{KERNEL_IMPORT_PLACEHOLDER}}
21
+
22
+ # {{UTILITY_FUNCTIONS_PLACEHOLDER}}
23
+
24
+ assert imported_kernel_function is not None, "imported_kernel_function is missing"
25
+
26
+ KERNEL_NAME = "{{KERNEL_NAME_PLACEHOLDER}}"
27
+ REPRO_CONTEXT_FILE_NAME = "{{JSON_FILE_NAME_PLACEHOLDER}}"
28
+
29
+
30
+ def _get_launch_kernel_args() -> Tuple[Tuple[int], Dict[str, Any]]:
31
+ script_dir = Path(__file__).resolve().parent # noqa: F821
32
+ json_file = script_dir / REPRO_CONTEXT_FILE_NAME
33
+
34
+ grid, args_dict = create_args_from_json_file(json_file) # noqa: F821, F841
35
+
36
+ print("Recorded kernel arguments dictionary:")
37
+ for name, arg in args_dict.items():
38
+ if isinstance(arg, torch.Tensor):
39
+ print(
40
+ f" {name}: Tensor: {arg.shape} {arg.dtype} stride: {arg.stride()}, is_contiguous: {arg.is_contiguous()}"
41
+ )
42
+ else:
43
+ print(f" {name}: {arg}")
44
+ print(f"Grid: {grid}")
45
+
46
+ return tuple(grid), args_dict
47
+
48
+
49
+ grid, args_dict = _get_launch_kernel_args()
50
+
51
+
52
+ def _launch_kernel(grid: tuple[int], args_dict: dict[str, Any]):
53
+ try:
54
+ assert grid is not None
55
+ assert args_dict is not None
56
+
57
+ # {{KERNEL_INVOCATION_PLACEHOLDER}}
58
+
59
+ except Exception as e:
60
+ print(f"Error: {e}")
61
+ print("Failed to launch kernel!")
62
+
63
+
64
+ # HACK: @register_x_val doesn't allow us to pass `operator_name`` as a parameter
65
+ tensor_args = {k: v for k, v in args_dict.items() if isinstance(v, torch.Tensor)}
66
+ x_vals_label = ", ".join(tensor_args.keys())
67
+ REGISTERED_X_VALS[KERNEL_NAME] = x_vals_label
68
+
69
+
70
+ class Operator(BenchmarkOperator):
71
+ @register_benchmark(operator_name=KERNEL_NAME)
72
+ def run_kernel(self, grid, args_dict):
73
+ return lambda: _launch_kernel(grid, args_dict)
74
+
75
+ def get_input_iter(self):
76
+ yield {"grid": grid, "args_dict": args_dict}
77
+
78
+ def get_x_val(self, example_inputs):
79
+ tensors_shapes = [
80
+ tuple(v.shape)
81
+ for v in example_inputs["args_dict"].values()
82
+ if isinstance(v, torch.Tensor)
83
+ ]
84
+ return tuple(tensors_shapes)
85
+
86
+
87
+ if __name__ == "__main__":
88
+ print("do_benchmark...")
89
+
90
+ args = [
91
+ "--benchmark-name",
92
+ KERNEL_NAME,
93
+ ]
94
+
95
+ from tritonbench.utils.parser import get_parser
96
+
97
+ parser = get_parser(args)
98
+ tb_args, extra_args = parser.parse_known_args(args)
99
+ bench = Operator(tb_args, extra_args)
100
+ bench.run()
101
+
102
+ print(bench.output)
103
+ print("Benchmark completed successfully!")
@@ -327,7 +327,7 @@ def _create_arg_from_info(arg_info):
327
327
  return None
328
328
 
329
329
 
330
- def determine_output_paths(out_dir: str, kernel_name: str):
330
+ def determine_output_paths(out_dir: str, kernel_name: str, template: str):
331
331
  """
332
332
  Determine output file paths for reproducer script and context data.
333
333
 
@@ -342,7 +342,12 @@ def determine_output_paths(out_dir: str, kernel_name: str):
342
342
  output_directory = Path(out_dir) / kernel_name
343
343
  output_directory.mkdir(parents=True, exist_ok=True)
344
344
 
345
- out_py_path = output_directory / f"repro_{timestamp}.py"
345
+ filename_parts = ["repro"]
346
+ if template != "example":
347
+ filename_parts.append(template.replace(".", "_"))
348
+ filename_parts.append(timestamp)
349
+ filename = "_".join(filename_parts) + ".py"
350
+ out_py_path = output_directory / filename
346
351
  temp_json_path = output_directory / f"repro_context_{timestamp}.json"
347
352
 
348
353
  return out_py_path, temp_json_path
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tritonparse
3
- Version: 0.3.1.dev20251028071524
3
+ Version: 0.3.1.dev20251029071541
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
@@ -62,6 +62,7 @@ tritonparse/reproducer/ingestion/ndjson.py
62
62
  tritonparse/reproducer/templates/__init__.py
63
63
  tritonparse/reproducer/templates/example.py
64
64
  tritonparse/reproducer/templates/loader.py
65
+ tritonparse/reproducer/templates/tritonbench.py
65
66
  tritonparse/tools/__init__.py
66
67
  tritonparse/tools/decompress_bin_ndjson.py
67
68
  tritonparse/tools/disasm.py
@@ -92,6 +93,7 @@ website/src/assets/react.svg
92
93
  website/src/components/ArgumentViewer.tsx
93
94
  website/src/components/Callstack.tsx
94
95
  website/src/components/CodeComparisonView.tsx
96
+ website/src/components/CodeViewer.css
95
97
  website/src/components/CodeViewer.tsx
96
98
  website/src/components/CompilationInfo.tsx
97
99
  website/src/components/CopyCodeButton.tsx
@@ -0,0 +1,43 @@
1
+ /* Location definition line styles */
2
+ .loc-definition-line {
3
+ background-color: rgba(255, 215, 0, 0.1) !important;
4
+ border-left: 3px solid #ffd700 !important;
5
+ }
6
+
7
+ .loc-definition-line:hover {
8
+ background-color: rgba(255, 215, 0, 0.2) !important;
9
+ }
10
+
11
+ /* Alias name badge */
12
+ .alias-badge {
13
+ display: inline-block;
14
+ margin-left: 8px;
15
+ padding: 2px 6px;
16
+ background: #e3f2fd;
17
+ border-radius: 3px;
18
+ font-size: 0.85em;
19
+ color: #1976d2;
20
+ font-weight: 500;
21
+ vertical-align: middle;
22
+ }
23
+
24
+ .alias-badge:hover {
25
+ background: #bbdefb;
26
+ }
27
+
28
+ /* Dark theme adjustments */
29
+ [data-theme="dark"] .alias-badge {
30
+ background: #1565c0;
31
+ color: #e3f2fd;
32
+ }
33
+
34
+ [data-theme="dark"] .alias-badge:hover {
35
+ background: #1976d2;
36
+ }
37
+
38
+ /* Loc ID indicator */
39
+ .loc-id-indicator {
40
+ opacity: 0.6;
41
+ font-size: 0.8em;
42
+ margin-left: 4px;
43
+ }
@@ -5,6 +5,7 @@ import {
5
5
  oneDark,
6
6
  } from "react-syntax-highlighter/dist/esm/styles/prism";
7
7
  import type { SourceMapping } from "../utils/dataLoader";
8
+ import "./CodeViewer.css";
8
9
 
9
10
  // Import language support
10
11
  import llvm from 'react-syntax-highlighter/dist/esm/languages/prism/llvm';
@@ -207,6 +208,7 @@ const BasicCodeViewer: React.FC<CodeViewerProps> = ({
207
208
  highlightedLines = [],
208
209
  onLineClick,
209
210
  viewerId,
211
+ sourceMapping,
210
212
  }) => {
211
213
  const containerRef = useRef<HTMLDivElement>(null);
212
214
  const lines = splitIntoLines(code);
@@ -239,6 +241,21 @@ const BasicCodeViewer: React.FC<CodeViewerProps> = ({
239
241
  {lines.map((line, index) => {
240
242
  const lineNumber = index + 1;
241
243
  const isHighlighted = highlightedLines.includes(lineNumber);
244
+ const mapping = sourceMapping?.[lineNumber.toString()];
245
+ const isLocDef = mapping?.type === "loc_def" || mapping?.kind === "loc_def";
246
+
247
+ // Build tooltip text for loc definitions
248
+ let tooltipText = "";
249
+ if (isLocDef && mapping) {
250
+ tooltipText = `Location definition`;
251
+ if (mapping.alias_name) {
252
+ tooltipText += `: ${mapping.alias_name}`;
253
+ }
254
+ if (mapping.alias_of !== undefined) {
255
+ const targetKey = mapping.alias_of === "" ? "#loc" : `#loc${mapping.alias_of}`;
256
+ tooltipText += ` → ${targetKey}`;
257
+ }
258
+ }
242
259
 
243
260
  return (
244
261
  <div
@@ -255,7 +272,11 @@ const BasicCodeViewer: React.FC<CodeViewerProps> = ({
255
272
  }}
256
273
  onClick={() => handleLineClick(lineNumber)}
257
274
  data-line-number={lineNumber}
258
- className={isHighlighted ? 'highlighted-line' : ''}
275
+ className={`${isHighlighted ? 'highlighted-line' : ''} ${isLocDef ? 'loc-definition-line' : ''}`}
276
+ title={tooltipText}
277
+ data-loc-id={mapping?.loc_id}
278
+ data-alias-name={mapping?.alias_name}
279
+ data-loc-def={isLocDef ? "true" : undefined}
259
280
  >
260
281
  <span style={{
261
282
  position: "absolute",
@@ -290,6 +311,7 @@ const LargeFileViewer: React.FC<CodeViewerProps> = ({
290
311
  highlightedLines = [],
291
312
  onLineClick,
292
313
  viewerId,
314
+ sourceMapping,
293
315
  }) => {
294
316
  const containerRef = useRef<HTMLDivElement>(null);
295
317
  const [visibleRange, setVisibleRange] = useState({ start: 0, end: 50 });
@@ -421,6 +443,8 @@ const LargeFileViewer: React.FC<CodeViewerProps> = ({
421
443
  lineProps={(lineNumber) => {
422
444
  // Adjust line number based on visible range
423
445
  const actualLine = lineNumber + visibleRange.start;
446
+ const mapping = sourceMapping?.[actualLine.toString()];
447
+ const isLocDef = mapping?.type === "loc_def" || mapping?.kind === "loc_def";
424
448
 
425
449
  // Create styles for the line
426
450
  const style: React.CSSProperties = {
@@ -440,11 +464,28 @@ const LargeFileViewer: React.FC<CodeViewerProps> = ({
440
464
 
441
465
  }
442
466
 
467
+ // Build tooltip text for loc definitions
468
+ let tooltipText = "";
469
+ if (isLocDef && mapping) {
470
+ tooltipText = `Location definition`;
471
+ if (mapping.alias_name) {
472
+ tooltipText += `: ${mapping.alias_name}`;
473
+ }
474
+ if (mapping.alias_of !== undefined) {
475
+ const targetKey = mapping.alias_of === "" ? "#loc" : `#loc${mapping.alias_of}`;
476
+ tooltipText += ` → ${targetKey}`;
477
+ }
478
+ }
479
+
443
480
  return {
444
481
  style,
445
482
  onClick: () => handleLineClick(actualLine),
446
483
  'data-line-number': actualLine,
447
- className: isHighlighted ? 'highlighted-line' : '',
484
+ className: `${isHighlighted ? 'highlighted-line' : ''} ${isLocDef ? 'loc-definition-line' : ''}`,
485
+ title: tooltipText,
486
+ 'data-loc-id': mapping?.loc_id,
487
+ 'data-alias-name': mapping?.alias_name,
488
+ 'data-loc-def': isLocDef ? "true" : undefined,
448
489
  };
449
490
  }}
450
491
  customStyle={{
@@ -474,6 +515,7 @@ const StandardCodeViewer: React.FC<CodeViewerProps> = ({
474
515
  highlightedLines = [],
475
516
  onLineClick,
476
517
  viewerId,
518
+ sourceMapping,
477
519
  }) => {
478
520
  const containerRef = useRef<HTMLDivElement>(null);
479
521
 
@@ -537,6 +579,22 @@ const StandardCodeViewer: React.FC<CodeViewerProps> = ({
537
579
  wrapLines={true}
538
580
  lineProps={(lineNumber) => {
539
581
  const isHighlighted = highlightedLines.includes(lineNumber);
582
+ const mapping = sourceMapping?.[lineNumber.toString()];
583
+ const isLocDef = mapping?.type === "loc_def" || mapping?.kind === "loc_def";
584
+
585
+ // Build tooltip text for loc definitions
586
+ let tooltipText = "";
587
+ if (isLocDef && mapping) {
588
+ tooltipText = `Location definition`;
589
+ if (mapping.alias_name) {
590
+ tooltipText += `: ${mapping.alias_name}`;
591
+ }
592
+ if (mapping.alias_of !== undefined) {
593
+ const targetKey = mapping.alias_of === "" ? "#loc" : `#loc${mapping.alias_of}`;
594
+ tooltipText += ` → ${targetKey}`;
595
+ }
596
+ }
597
+
540
598
  return {
541
599
  style: {
542
600
  backgroundColor: isHighlighted
@@ -549,7 +607,11 @@ const StandardCodeViewer: React.FC<CodeViewerProps> = ({
549
607
  },
550
608
  onClick: () => handleLineClick(lineNumber),
551
609
  "data-line-number": lineNumber,
552
- className: isHighlighted ? 'highlighted-line' : '',
610
+ className: `${isHighlighted ? 'highlighted-line' : ''} ${isLocDef ? 'loc-definition-line' : ''}`,
611
+ title: tooltipText,
612
+ "data-loc-id": mapping?.loc_id,
613
+ "data-alias-name": mapping?.alias_name,
614
+ "data-loc-def": isLocDef ? "true" : undefined,
553
615
  };
554
616
  }}
555
617
  >
@@ -17,6 +17,12 @@ export interface SourceMapping {
17
17
  ttgir_lines?: number[]; // Array of corresponding TTGIR lines
18
18
  llir_lines?: number[]; // Array of corresponding LLIR lines
19
19
  amdgcn_lines?: number[]; // Array of corresponding AMDGCN lines
20
+ // New fields for location alias support
21
+ type?: string; // Type of mapping entry, e.g., "loc_def" for loc definition lines
22
+ kind?: string; // Deprecated alias for type, kept for backward compatibility
23
+ loc_id?: string; // The #loc identifier (e.g., "13" for #loc13)
24
+ alias_name?: string; // Name of the alias (e.g., "x_ptr" in #loc13 = loc("x_ptr"(#loc)))
25
+ alias_of?: string; // The target #loc this alias points to
20
26
  }
21
27
 
22
28
  /**
@@ -49,19 +49,19 @@ const IRAnalysis: React.FC<IRAnalysisProps> = ({ kernels, selectedKernel }) => {
49
49
  {ttgir_info && (
50
50
  <>
51
51
  <div className="flex flex-col">
52
- <span className="text-sm font-medium text-gray-500">Tiled Buffer Load Count</span>
52
+ <span className="text-sm font-medium text-gray-500">Tiled Global Load Count</span>
53
53
  <span className="font-mono text-sm break-words">{getCount(ttgir_info, "tt.load_count")}</span>
54
54
  </div>
55
55
  <div className="flex flex-col">
56
- <span className="text-sm font-medium text-gray-500">Tiled Buffer Store Count</span>
56
+ <span className="text-sm font-medium text-gray-500">Tiled Global Store Count</span>
57
57
  <span className="font-mono text-sm break-words">{getCount(ttgir_info, "tt.store_count")}</span>
58
58
  </div>
59
59
  <div className="flex flex-col">
60
- <span className="text-sm font-medium text-gray-500">Tiled Global Load Count</span>
60
+ <span className="text-sm font-medium text-gray-500">Tiled Buffer Load Count</span>
61
61
  <span className="font-mono text-sm break-words">{getCount(ttgir_info, "amdgpu.buffer_load_count")}</span>
62
62
  </div>
63
63
  <div className="flex flex-col">
64
- <span className="text-sm font-medium text-gray-500">Tiled Global Store Count</span>
64
+ <span className="text-sm font-medium text-gray-500">Tiled Buffer Store Count</span>
65
65
  <span className="font-mono text-sm break-words">{getCount(ttgir_info, "amdgpu.buffer_store_count")}</span>
66
66
  </div>
67
67
  </>