o wZh@s,ddlZddlZddlZddlmZdedefddZdedefddZdedefd d Zd e ed e ede efd dZ dede ee eeffdefddZ dededefddZedfdedefddZedkr ejddZejdedd ejd!eedd"d#eZeejejZdSdS)$N)Path source_codereturncCs tdd|}tdd|}|S)Nz(\n.+\s\'\'\'\n) z (\'\'\'\,.+))resub)rZ remove_headZ remove_tailrL/var/www/auris/lib/python3.10/site-packages/torch/utils/_get_clean_triton.py"remove_triton_function_declarationsr cCs.t|dd}t|dd}t|dd}|S)Nzasync_compile = AsyncCompile()zasync_compile.wait(globals())zdel async_compile)strreplace)rZremove_top_levelZremove_compileZ remove_delrrr remove_async_compilesrcCsxd}d}ddt||tjD}t|D]$\}}|||}|dkr9|d|d|||t|d}qq|S)Nz/(\w+)\s*=\s*async_compile\.triton\('triton_',\sz def triton_cSsg|] }||dfqS))endgroup).0matchrrr sz"rename_kernels..zdef )rfinditerDOTALLreversedfindlen)rpatternZtriton_kernel_declmatchesZ end_indexZcaptured_stringZ index_of_Brrr rename_kernelss   roriginal_params new_paramscCs.tt|D]}||dkr||||<q|S)NT)ranger)rridxrrr merge_params/s   r#originalkernel_to_paramscs<d}dtffdd }t|||}tjdd|tjd}|S)Nz(\w+)\.run\((.*)\)rcsV|d}|d}|\}}t|d|d}|d|dd|d}|S)Nrz, [z]())rr#splitjoin)r func_nameparamsrgridZ new_stringr%rr r <s   z"add_launch_params..replacez"@triton_heuristics[^@]*@triton.jitz @triton.jit)flags)r rrr)r$r%rr Z transformedZremove_inductor_wrappersrr.r add_launch_params6s r0input_filenameoutput_filenamec Cst| }|}Wdn1swY|}d|vr"tdt|}t|}|d}tj|s@td|d|dt| }|}Wdn1sSwYdd|D}d d |D} t || }t|d }| |Wd|S1swY|S) Nz def triton_(z]Need to run original Pytorch code generating kernels with TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1z.launch_paramszMissing z1. Run `TORCHINDUCTOR_DUMP_LAUNCH_PARAMS=1 python z first.cSsg|]}|dqS)|)r))rirrr risz process_file..cSs(i|]\}}}|||fqSr)strip)rabcrrr js(z process_file..w) openread RuntimeErrorr rospathexists readlinesr0write) r1r2filerZtransformed_codeZlaunch_params_filenamefZlaunch_params_metaZ split_paramsZkernel_args_gridrrr process_fileSs6          rEztriton_only_repro.py input_path output_pathcCstt|t|S)zRun experiments and output results to file Args: input_path (Optional[Path]): Path to inductor generated output codede output_path (Optional[Path]): Path to write out the new python file )rEr )rFrGrrr get_clean_tritonrs rH__main__z=Clean Inductor generated code to remove Inductor dependencies) descriptionz&Path to inductor generated output code)typehelpz --output_pathz)Path to write out the clean triton output)rKdefaultrL)argparser>rpathlibrr r rrlistr#dicttupler0rErH__name__ArgumentParserparser add_argument parse_argsargsrFrGresultrrrr sN "