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

Merge pull request #23369 from sergachev:cleanup_nsys_converter

Browse Source 
PiperOrigin-RevId: 671801177
This commit is contained in:
jax authors 2024-09-06 09:54:59 -07:00
commit 878b6b5743
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/tools/pgo_nsys_converter.py
View File

@ -28,8 +28,7 @@ if __name__ == '__main__':
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")
parser.add_argument("--pgle_output_path", type=str, help="output file", default="/opt/paxml/workspace/lhs_pbtxt/temp.pbtxt")
args = parser.parse_args()