mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

respond to reviewer's comments

Browse Source
This commit is contained in:
Abhinav Goel 2024-03-04 11:42:01 -08:00 committed by GitHub
parent 5a732ad89f
commit 2480ca383e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/tools/pgo_nsys_converter.py
View File

@ -7,9 +7,11 @@ import os
import shutil
import subprocess
print("Script to convert NVIDIA Nsys Profiles to the .pbtxt format. This format is readable by XLA's Profile Guided Latency Estimator. Usage: pgo_nsys_converter.py --profile_path <path the nsys profile> --pgle_output_path <path to output .pbtxt>")
nsys_path = shutil.which("nsys")
parser = argparse.ArgumentParser(description='Tool to sweep for optimal collective combiner threshold')
parser = argparse.ArgumentParser(description='Tool to convert NVIDIA Nsys Profiles to the .pbtxt format')
parser.add_argument("--profile_path", type=str, help="path to nsys profile")
parser.add_argument("--post_process", help="post process pbtxt to get minimum cost value for each instruction", action="store_true")
parser.add_argument("--pgle_output_path", type=str, help="output directory", default="/opt/paxml/workspace/lhs_pbtxt/temp.pbtxt")