update reproducer template

Summary: This diff contains minor fixes and improvements to the reproducer template system.

Reviewed By: FindHao

Differential Revision: D85655006

fbshipit-source-id: 6f827073820daf1916bc66a815c7020565f84198

Author: wychi

tritonparse/reproducer/orchestrator.py

@@ -23,7 +23,7 @@ def reproduce(
     template: str,
     replacer: Optional[PlaceholderReplacer] = None,
     kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
-) -> dict[str, Path]:
+) -> dict[str, str]:
     """
     Generate a reproducer script from NDJSON trace file.
 

tritonparse/reproducer/templates/example.py

@@ -14,7 +14,7 @@
 # {{UTILITY_FUNCTIONS_PLACEHOLDER}}
 
 
-if __name__ == "__main__":
+def launch_kernel():
     script_dir = Path(__file__).resolve().parent  # noqa: F821
     json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
     grid, args_dict = create_args_from_json_file(str(json_file))  # noqa: F821
@@ -28,3 +28,7 @@
 
     torch.cuda.synchronize()
     print("Kernel execution finished.")
+
+
+if __name__ == "__main__":
+    launch_kernel()