9j"9TUddlmZddlZddlZddlZddlZddlZddlmZddlm Z ddl m Z ddl m Z ddlmZddlmZe rdd lmZdd lmZdd lmZdadadagad ed <gZded<dadaejGddZ ga!ded<da"da#da$da%ded<da&da'dd*Z/d?d+Z0d@d,Z1dAd-Z2dBd.Z3d?d/Z4dCd0Z5dDd1Z6 dEd2Z7 dFd3Z8dy)L) annotationsN) dataclass) lru_cache) TYPE_CHECKING)config)get_benchmark_name) OrderedSet)Callable)Config)BaseSchedulerNodez#list[tuple[BaseSchedulerNode, int]]nodes_num_elemz%list[tuple[BaseSchedulerNode, float]] node_runtimesc&eZdZUded<dZded<y)CppOuterLoopFusedCountintinner_kernel_numberrlocal_buffer_numberN)__name__ __module__ __qualname____annotations__rW/media/conek/DATA/Code/OCR/venv/lib/python3.12/site-packages/torch/_inductor/metrics.pyrr+s  rrzlist[CppOuterLoopFusedCount]!cpp_outer_loop_fused_inner_countsrnum_auto_chunkingcdadadatj t j dadatj da da da da da day)Nr)generated_kernel_countgenerated_cpp_vec_kernel_countnum_bytes_accessedr clearrir_nodes_pre_fusioncpp_to_dtype_countrnum_comprehensive_padding)num_matches_for_scatter_upon_const_tensornum_loop_reorderingparallel_reduction_countcodegen_mix_order_reductionrrrrresetr)Asj%&"%++- !01- "#rcNeZdZUdZded<ded<ded<ded<ded<ded<y ) CachedMetricsDeltasz] The subset of metrics we want update across cache hits, e.g., the FxGraphCache. rrrr"r#r r%N)rrr__doc__rrrrr+r+_s-  $''/22rr+cntjtDcgc]}|jc}Scc}wN) dataclassesfieldsr+name)fields rget_metric_fieldsr3ns&$/$6$67J$K L5EJJ LL Ls2c2eZdZdZddZddZeddZy) CachedMetricsHelperz A helper class to help calculate and apply counter deltas for those metrics we want to save with cache entries (e.g., FxGraphCache) and apply on a cache hit. cdi|_tD]}t||j|<yr.)cached_metricsr3globals)selfmetrics r__init__zCachedMetricsHelper.__init__ys3 ') staticmethodrBrrrr5r5rs% < 488rr5zdict[str, MetricTable]REGISTERED_METRIC_TABLEScbeZdZUded<ded<dZded<ddZdd Zdd Zdd Ze dd Z y ) MetricTablestr table_name list[str] column_namesrrnum_rows_addedcF|jtvry|}t|jt|k(s(Jt|jdt|t |jt |j k(s6Jt |jdt |j t }|g|jDcgc]}|| c}z}td|DsJ|j|ycc}w)Nz v.s. c 3\K|]$}t|tttdf&ywr.) isinstancerJfloattype).0is r z&MetricTable.add_row..s"Hq:a#ud4j!9:Hs*,) rKenabled_metric_tableslenrMr keysrall _write_row)r9row_fnrow_dictbn column_namerows radd_rowzMetricTable.add_rows ??"7"9 9 84$$%X6 4$$%&fS]O < 6$++, 8==?0KK $++,-VJx}}4O3P Q K !dt?P?PQ h{+QQHCHHHH Rs' Dc"d|jdS)N metric_table_z.csv)rK)r9s routput_filenamezMetricTable.output_filenamest/t44rc|j}t|d5}tj|d}|j dg|j zdddy#1swYyxYw)Nw lineterminator model_name)rdopencsvwriterwriterowrM)r9filenamefdrms r write_headerzMetricTable.write_headers^'') (C  @BZZ48F OO\NT->->> ? @ @ @s 7AA&c|j}|jdk(r/tjj |s|j |xjdz c_t |D]&\}}t|tr|d}n|d}n|}|||<(t|d5}tj|d}|j|dddy#1swYyxYw)Nrz.6fargrh) rdrNospathexistsrq enumeraterQrRrkrlrmrn)r9r`roidxorig_valnew_valrprms rr[zMetricTable._write_rows'')   ! #BGGNN8,D     q &s^ MC(E*%cN!"CH (C  !BZZ48F OOC  ! ! !s $)CCc.t||}|t|<yr.)rIrG)r1rMtables rregister_tablezMetricTable.register_tablesD,/). &rN)r\z+Callable[[], dict[str, str | float | None]]rDrE)rDrJrC)r`zlist[str | float | None]rDrE)r1rJrMrLrDrE) rrrrrNrardrqr[rFrrrrrIrIs@ONC$5@ !(//rrI slow_fusion) kernel1_pathkernel1_latency kernel2_pathkernel2_latencyfused_kernel_pathfused_kernel_latencyslow_down_ratio graph_stats)graph_idnum_nodes_before_fusionnum_nodes_after_fusionpersistent_red_perf) kernel0_pathrr kernel3_pathkernel0_latencyrrkernel3_latency size_hintsreduction_hint'fusion_failure_due_to_indexing_mismatch)pre_grad_graph_idpost_grad_graph_id node1_name node2_namenode1_debug_strnode2_debug_strcommon_buffer_namesfailure_reasonkernel_metadata kernel_name kernel_pathkernel_categoryrr line_of_codenum_load num_store num_for_loopnum_atomic_addnum_argsxnumelynumelrnumelkernel_args_num_gbcddlm}ddlm}|j |}||}t j |jjS)z The kernel_module_code is the python module that contains kernel function code. kernel function is the proper triton kernel function annotated with @triton.jit rs) PyCodeCache)get_triton_kernel) codecacherwrapper_benchmarkrloadinspect getsourcefn)kernel_module_coderrmodkernels r_parse_kernel_fn_coders@ '4   - .C s #F   VYY\\ **rc4t|jS)zJ Return the line of code for the kernel excluding the decorators. )rX splitlines)proper_kernel_fn_codes r_parse_kernel_line_of_coder.s $//1 22rcn|dk(rytjd|}|sJd|jdS)Nforeachzsize_hints=(\[[0-9, ]*\]),zsize_hints missing!rsresearchgroup)rrms r_parse_size_hintsr5s;)# /1CDA ###1 771:rcl|dvrytjd|}|sJd|jdS)N) reductionpersistent_reductionz$reduction_hint=ReductionHint\.(\w*),z/reduction_hint not found in kernel source code!rsr)rrrs r_parse_reduction_hintr>s<CC 9;MNA ???1 771:rc$|j|Sr.)count)rpatterns r_count_patternrFs & &w //rc|jd}|jdsJ|jd}|jd}||dz|}|jd}t |S)Nrdef (z):rs,)r startswithindexsplitrX)rdef_line start_idxend_idxdecl_csvcompss r _count_argsrJsm$//1!4H   v && &s#InnT"G A 0H NN3 E u:rc.|jd}||dS)z Skip decorators. rN)r)kernel_fn_code start_poss r_parse_proper_kernel_fn_coderTs!$$V,I )* %%rcntj|d|}|rt|jdSy)Nz = ([\d]+)rs)rrrr)rnumel_arg_namers r _parse_numelr\s3 ^$K02GHA1771:rcjtjd|}|rt|jdS y)z inductor meta looks like: inductor_meta={... 'mutated_arg_names': [], 'no_x_dim': False, 'kernel_num_gb': 2.0}, z.kernel_num_gb.:\s*([0-9.]+)rsN)rrrRr)rrrs r_parse_kernel_args_num_gbrds5 1>BAQWWQZ   rc  ddlm}||t|t| t |t t tdj fdy)z An utility to log kernel metadata. We may parse metadata from kernel source code here. It's fine to parse the generated kernel code here since the logging is disabled by default. It would hurt compilation time. rs)"get_kernel_category_by_source_coderctdtdtdtdttdtdtdtdS) Nztl.loadztl.storezfor z tl.atomic_addrrrr)rrrr)rrkernel_line_of_coderrrrrsrz%log_kernel_metadata..s&&.$,/&'N8H55JK&'//  rctjD]f\}}|tvs|j}tj j |rt j||jhy)z Purge the old log file at the beginning when the benchmark script runs. Should do it in the parent process rather than the child processes running each individual model. N) rGitemsrWrdrvrwrxunlinkrq)r1r~ros rpurge_old_log_filesrsb 0557! e (* *,,.Hww~~h' (#    !rc4ttjSr.)enabled_metric_tables_implrrWrrrrWrWs %f&B&B CCrct}|jdD]9}|j}|s|tvs Jd|d|j |;|S)NrzMetric table name z is not registered)r rstriprGadd) config_strenabledr1s rrrsl)|G  %zz| //  &8 9 /  D Nrc|tvSr.)rWr1s ris_metric_table_enabledrs (* **rc:|tvs Jd|dt|S)Nz Metric table z is not defined)rGrs rrrs* + +R}TF/-RR + #D ))rkernel_autotunerr triton_config latency_mscLtdjfdy)Nrc$tdS)Nr)rJ)rrrlatencysrrz,log_kernel_autotune_result..s&& [!  r)rra)rrrrs````rlog_kernel_autotune_resultrs&'// rrC)rDrL)rrJrDrJ)rrJrDr)rrJrrJrD str | None)rrJrrJrDr)rrJrrJrDr)rrJrDrJ)rrJrrJrDz int | None)rrJrrJrDz float | None)rrJrrJrrJrDrE)rDOrderedSet[str])rrJrDr)r1rJrDbool)r1rJrDrI) rrJrrJrr rrRrDrE)? __future__rrlr/rrvrr functoolsrtypingrtorch._inductorrtorch._inductor.utilsrtorch.utils._ordered_setr collections.abcr %torch._inductor.runtime.triton_compatr torch._inductor.schedulerr rrr r rrr"r#rrr$r%r&rr'r(r)r+r3r5rGrIrrrrrrrrrrrrrWrrrrrrrr s}"  ! "4/(<;!"   8: 49 !!! CE!#?D,-)3<  3 3  3M8824605 8/8/ 8/v     " -   6 + 30&*-*))#&)