jpDddlZddlZddlZddlZddlZddlZddlZddlZddlm Z ddl m Z m Z m Z ddlmZddlmZdZdZd Ze d Zd Zd"dZd"dZd#dZd#dZd$dZGddZGddZdZd%dZd%dZ dZ!e d&d Z"d%d!Z#dS)'N)contextmanager)AnyDictList)language)runtimecd|}dddd|zdg}tj|}|tjjd}d|D}|S)N, nvidia-smi-i0z --query-gpu=z--format=csv,noheader,nounitsc,g|]}t|S)int).0xs s/home/lighthouse/badge-speaker-benchmark/.benchmark-venv-community/lib64/python3.11/site-packages/triton/testing.py znvsmi..s   a3q66   )join subprocess check_outputdecodesysstdoutencodingsplit)attrscmdoutrets rnvsmir#ss HHUOOE sNU$:<[ \C  !# & &C **SZ( ) ) / / 4 4C  3   C Jrclttfdfd|DS)Ncd|cxkrdksntd|dz z}tj|}tj|}||z }d|z |z||zzS)Nrrz%Quantiles must be in the range [0, 1]) ValueErrormathfloorceil)qpointloweruppertans r get_quantilez_quantile..get_quantile s}Q ! DEE EQU  5!! %   EMA5!A%L00rc&g|] }|Srr)rqir1s rrz_quantile..)s# ) ) )LL   ) ) )r)lensorted)r/r*r1r0s` @@r _quantiler6sU AAq A111111 * ) ) )q ) ) ))rc0|-t||}t|dkr|d}|S|dkr|S|dkrt|S|dkrt|S|dkrt j|S|dkrt j|SdS)Nrrallminmaxmeanmedian)r6r4r9r: statisticsr;r<)times quantiles return_moder"s r_summarize_statisticsrA,sy)) s88q==a&C e   5zz   5zz   u%%%   ''' ! rc#Kddlm}tjdt jj5}|tj |ddd} ||fV|| |n#|| |wwxYwddddS#1swxYwYdS)Nrztriton-bench-proton-)prefixprofileshadowtree)contextdata) triton.profilerprofilertempfileTemporaryDirectoryuuiduuid4hexstartospathrfinalize)protontmpdirsessions r_proton_bench_sessionrW>s$$$$$$  $,U4:<>W],^^ )'/ ! ! !"((("((((# ))))))))))))))))))s)7B7/B 5B7 B''B77B;>B;czgfdfd|D] }|dtDS)Nc|dg}t|dkr-|didddz Stfd|DS)Nchildrenrmetricsz time (ns).Ac3.K|]}|VdSNr)rchildkernel_time_mss r zF_collect_proton_scope_times..kernel_time_ms..Rs-??U>>%((??????r)getr4sum)noderZr`s rr`z3_collect_proton_scope_times..kernel_time_msNss88J++ x==A  88Ir**..{A>>D D????h??????rc"|didd}|r*|}|dkr||fdS|dgD] }|dS)NframenamerrZ)rb startswithappend)rdrgtime_msr_r`rC scope_timesvisits rrmz*_collect_proton_scope_times..visitTsxx$$((44 ??6 " " $nT**G{{""D'?333 FXXj"--  E E%LLLL  rcg|]\}}|Srr)r_times rrz/_collect_proton_scope_times..`s 4 4 4WQD 4 4 4r)r5)databaserCrdr`rlrms ` @@@r_collect_proton_scope_timesrrKsK@@@@@  d 4 4{ 3 3 4 4 44rr;cddl}|dvsJ|j|j5||5|D]2}||dd|_3|jd}|jd}|tdD] } | ||j | |dz } | dkrd} n tdt|| z } |j} |j| 5t| D]} | |D] }d|_ | dddn #1swxYwY|j g} d }t|D]} |jd}|jd}|| ||j | | || z gz } t#| ||cdddS#1swxYwYdS) a Benchmark the runtime of the provided function. :param fn: Function to benchmark :type fn: Callable :param rep: Repetition time (in ms) :type rep: int :param grad_to_none: Reset the gradient of the provided tensor to None :type grad_to_none: torch.tensor, optional :param return_mode: The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean". :type return_mode: str rNr9r:r;r<r8T enable_timingr )torchcudastreamStreamdetach_requires_grad_gradEventrecordrange synchronize elapsed_timer:r CUDAGraphgraphreplayrA)fnrep grad_to_noner?r@r{r start_event end_eventro estimate_msn_repeatgr" n_retriess rdo_bench_cudagraphrcs|LLL A A A A A   5:,,.. / /0B0B   #!     &&&j&&T&:: J$$4$88 q  A BDDDD    !..y99A= !  HH1c# "34455H J " " Z  a   8__  +)&&!%                       y!! D DA****>>K ((t(<.spppqQ]pppr)#r driveractiveget_current_targetbackend RuntimeErrorr{r|r}r~rrrrrrrrr:rrWget_empty_cache_for_benchmarkrrMrNrOr clear_cachescoper deactivaterrrHrbrA)rrrr?r@targetr{rrrrorrrTrVcacher scope_prefixir>rs @rdo_bench_cudagraph_protonrs& A A A A A ^ " 5 5 7 7F ~TUU ULLL   5:,,.. / /-D-D   #!     &&&j&&T&:: J$$4$88 q  A BDDDD    !..y99A= &!++44QC+q&>&>&>??                 J " " $ $ $I :y))AHHJJJJ &&(((!!'D!9999!!'D!9999pppp,G X_H`H`bn,o,opppE3 q q q q q q q q q q q q q q q6%UI{CC[-D-D-D-D-D-D-D-D-D-D-D-D-D-D-D-D-D-Ds:D#N4BN $AJ2 J J2JJ2 J!J2& N 2J6 6N 9J6 :N ?L0N 0M  7N  N4 N N4N N44N8;N8dc|dvsJtjj|tjj}d}d}|tdD]0} tjj ||1|| |dz } tdt|| z } tdt|| z } fdt| D}fdt| D}t| D] } | t| D]r} | |D] }d|_ tjj ||| ||| sd t||D}t|||S) a Benchmark the runtime of the provided function. By default, returns the mean runtime of :code:`fn` as a single float. :param fn: Function to benchmark :type fn: Callable :param warmup: Warmup time (in ms). Controls how long the function is run before timing begins. :type warmup: int :param rep: Repetition time (in ms). Controls the total duration of the timed runs. :type rep: int :param grad_to_none: Reset the gradient of the provided tensors to None before each run, to avoid gradient accumulation affecting timing. :type grad_to_none: list[torch.Tensor], optional :param quantiles: If provided, return these quantiles of the runtime distribution instead of a summary statistic. For example, ``[0.2, 0.5, 0.8]`` returns the 20th, 50th, and 80th percentile runtimes in ms. When set, ``return_mode`` is ignored. :type quantiles: list[float], optional :param return_mode: The summary statistic to return when ``quantiles`` is not set. ``"mean"`` and ``"median"`` return a single float. ``"min"`` and ``"max"`` return the fastest and slowest run respectively. ``"all"`` returns the raw list of per-run timings in ms. Default is ``"mean"``. :type return_mode: str :return: A single float by default, a list of floats if ``quantiles`` is provided, or a list of all per-run timings if ``return_mode="all"``. :rtype: float | list[float] ruTrvrxrc<g|]}dSTrvrrrdis rrzdo_bench..#s'IIIA288$8//IIIrc<g|]}dSrrrs rrzdo_bench..$s'GGG!--GGGrNc>g|]\}}||Sr)r)rses rrzdo_bench..8s( G G G41aQ^^A   G G Gr)r rrget_device_interfacerrrrrrrr:rrziprA)rwarmuprrr?r@rrrrorn_warmuprrrr>rs @rdo_benchrsn4 A A A A A   3 3 5 5BBDDDNN N ! ? ? A AE(((..Kt,,I 1XX  ))%000  NN**9559K1c&;.//00H1c# +,,--HIIIIxIIIKGGGGuXGGGI 8__    8__    #!  ))%000A ! NN G G3{I+F+F G G GE  ; ? ??rc|dvsJtjj}||tjj}|d}|d} |tdD]0} tjj ||1| || | dz } | dkrd} d} n@tdt|| z } tdt|| z } t| D] } | |dtjjd }t!5\}}|t#d  t| D]p}| |D] }d |_ tjj ||||d 5|d d d n #1swxYwYq|||d n#||d wxYwt+|j||}d d d n #1swxYwYt1|||S)a Benchmark the runtime of kernels invoked by the provided function using the Proton profiler. The measured runtime is generally more accurate than `do_bench` for short kernels that are affected by CPU overhead. Note that this function has several constraints compared to `do_bench`: - It does not measure GPU operations other than kernels (e.g., memory copies, synchronization, etc.). - It supports only AMD and NVIDIA GPUs. :param fn: Function to benchmark. :type fn: Callable :param warmup: Warmup time (in ms). :type warmup: int :param rep: Repetition time (in ms). :type rep: int :param grad_to_none: Reset the gradient of the provided tensor(s) to `None`. :type grad_to_none: torch.Tensor, optional :param quantiles: Performance percentiles to return in addition to the median. :type quantiles: list[float], optional :param return_mode: The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean". :type return_mode: str ruTrvrxrryrrrNrrr)r rrrrrrrrrrr:rrMrNrOrWrrrrrrrHrbrA)rrrrr?r@rrrrrorrrrrTrVrrr>s rdo_bench_protonr<s, A A A A A   3 3 5 5BBDDDNN N ! ? ? A AE(((..Kt,,I 1XX  ))%000  NN**9559Kaq#f{23344q#cK/0011 8__   NN0TZ\\-000L  T$5VW ?BCC C 68__  +)&&!%%11%888\\\":1":":":;;BDDD NN      g  5 5 5 5F  g  5 5 5 5+FKOOG,D,DlSS#TTTTTTTTTTTTTTT& ! ; ? ??sOK"&AJ I JI J!I "J:K"J+++K""K&)K&rhc ddl}ddl}t||js||}t||js||}|d}t |r||jn|}|d}t |r||jn|}t||jr\|j|jkr|}| }t||jr\|j|jkr|}| }|j dks |j dkr"|j ||||d|dS|||||st|d |d |d |d |d dS)a Asserts that two inputs are close within a certain tolerance. :param x: The first input. :type x: scala, list, numpy.ndarray, or torch.Tensor :param y: The second input. :type y: scala, list, numpy.ndarray, or torch.Tensor :param atol: The absolute tolerance. Default value is 1e-2. :type atol: float, optional :param rtol: The relative tolerance. Default value is 0. :type rtol: float, optional :param err_msg: The error message to use if the assertion fails. :type err_msg: str rNg{Gz?grT)atolrtol equal_nanerr_msg)rr z is not close to z (atol=z, rtol=))numpyr{ isinstanceTensortensorcallabledtypebfloat16floatcpudetachsizetestingassert_allcloseallcloseAssertionError)ryrrrnpr{s r assert_closersLLL a & & LLOO a & & LLOO |$TNN 444===D |$TNN 444===D!U\""% 7en $ $ A EEGGNN   " " $ $!U\""% 7en $ $ A EEGGNN   " " $ $ vzzQVaZZ ""1adW^"___ ;;q!$T; 2 2^\\!\\a\\\\UY\\\]]]^^rceZdZdZ ddeedeededeed eed ed eeefd ed ededefdZ dS) Benchmarkzk This class is used by the :code:`perf_report` function to generate line plots with a concise API. rhFNx_namesx_valsline_arg line_vals line_names plot_nameargsxlabelylabelx_logy_logc ||_||_| |_||_||_||_| |_| |_||_| |_ ||_ ||_ dS)aq Constructor. x_vals can be a list of scalars or a list of tuples/lists. If x_vals is a list of scalars and there are multiple x_names, all arguments will have the same value. If x_vals is a list of tuples/lists, each element should have the same length as x_names. :param x_names: Name of the arguments that should appear on the x axis of the plot. :type x_names: List[str] :param x_vals: List of values to use for the arguments in :code:`x_names`. :type x_vals: List[Any] :param line_arg: Argument name for which different values correspond to different lines in the plot. :type line_arg: str :param line_vals: List of values to use for the arguments in :code:`line_arg`. :type line_vals: List[Any] :param line_names: Label names for the different lines. :type line_names: List[str] :param plot_name: Name of the plot. :type plot_name: str :param args: Dictionary of keyword arguments to remain fixed throughout the benchmark. :type args: Dict[str, Any] :param xlabel: Label for the x axis of the plot. :type xlabel: str, optional :param ylabel: Label for the y axis of the plot. :type ylabel: str, optional :param x_log: Whether the x axis should be log scale. :type x_log: bool, optional :param y_log: Whether the y axis should be log scale. :type y_log: bool, optional :param styles: A list of tuples, where each tuple contains two elements: a color and a linestyle. :type styles: list[tuple[str, str]] N) rrrrrrrstylesrrrr) selfrrrrrrrrrrrrs r__init__zBenchmark.__init__s]^     "$    " r)rhrhFFN) __name__ __module__ __qualname____doc__rstrrrboolrrrrrrs;;c;S ; ; 9 ; I ;;38n;;;;;;;;;;rrc :eZdZdZ d dedededefdZd d Zd S)Markc"||_||_dSr^r benchmarks)rrrs rrz Mark.__init__s$rFbench save_path show_plots print_datac V #ddl}ddlm} ddl} fdjD} fdjD} fdjD} t j}| || z| z| z}jD]#t#t tfs#fd|D#t#t|kr"tdt|d#tt|#}ggg}}}jD]Q}|jdi|j|ij|} |\}}}n#t&$r |dd}}}YnwxYw||gz }||gz }||gz }Rt #|z|z|z|jt|<jr<| | }|d}t1t| | | D]\}\}}}||||}}jrj|dnd}jrj|d nd} |||||||| |sz|sT|t<}|t<}|||||d | | |!j"p||#j$|%j&rd nd|'j(rd nd|r| )|r6| *|j+,|jd||| z}|rA|j-d dkr0|j./\}!}"||"||!z |d<|r8tajdzta|1|r=|2|j+,|jdd|dd|S)Nrc*g|]}|djdS)z (rrrrrs rrzMark._run.. s,KKKQA00000KKKrc*g|]}|djdS)z-min (rrrs rrzMark._run.. ,NNN133EL333NNNrc*g|]}|djdS)z-max (rrrs rrzMark._run.. rr)columnscg|]}Srr)rrors rrzMark._run..s(((1Q(((rz Expected z values, got r)labelcolorlsg333333?)alpharloglinearz.pngDiff:z.csvz%.fF) float_formatindexr)3rQmatplotlib.pyplotpyplotpandasrlistr DataFramerrtupler4r&dictrrrrr TypeErrorlocrfiguresubplot enumeraterplotisnullr8astyper fill_betweenlegend set_xlabelr set_ylabelr set_xscaler set_yscalershowsavefigrRrshapertolistprint to_stringto_csv)$rrrrrdiff_colsave_precisionkwargsrQpltpd y_mean_labels y_min_labels y_max_labelsrdfx_argsrow_meanrow_minrow_maxrr"y_meany_miny_maxaxfirst_xr mean_label min_label max_labelcolstycol0col1rs$ ` @r_runz Mark._runs ''''''KKKK%:JKKK NNNNU=MNNN NNNNU=MNNN u}%% \\'M"9L"H<"W\ X X E EAa$// )(((((((1vvW%% !KS\\!K!K!K!KLLL#gq//**F)+RwgH_ # #dgVVV5>1*=VVvVV;+.(FE55 ;;;+.d5EFFF;VH$E7"E7""1gg07:WDBF3r77OO ? O JJLLLBajG9B3}VbdpCqCq9r9r V V55J 9!)}bmu,1LBel1oa((d,1LBel1oa((d7 R ^:SUXYYY||~~))++VELLNN4F4F4H4HV!LL//E!LL//EOOBwKTQTOUUU IIKKK MM%,1' 2 2 2 MM%, ' ' ' MM5;<%%H = = = MM5;<%%H = = =   O BGLLu4L4L4LMMNNN -' (  - q((**,,JD$DBtH,BvJ  " %/C' ( ( ( ",,.. ! ! !  # IIbgll9.F.F.FGGVl[iVlVlVl!  # # # s?EEErhc t|jt}|r|jgn|j}g} |D]&} ||j| |||fi|' |rt j|dtt j |dd5} | d|dt|D] } | d| j d!| ddddn #1swxYwYn#|rt j|dtt j |dd5} | d|dt|D] } | d| j d!| ddddw#1swxYwYwwxYw|r |r|d S|SdS) NT)exist_okz results.htmlwz z z r) rrrrjrArQmakedirsopenrRrwriter4r) rrrr return_dfr*has_single_benchr result_dfsrhtmls rrunzMark.runHs%doyAA*:Odo&&   3# a a!!)$)E9j*"_"_X^"_"_```` a 3 I5555"',,y.AA3GG34JJ/000!+,zperf_report..gsb*--rr)rwrappers` r perf_reportrQ`s.---G Nrcddlm}|+|j}|jj|d}|jj|d}||zdzdz dz }|S) z return DRAM bandwidth in GB/s rrNmem_clock_rate mem_bus_widthrr\)r rrrcurrent_deviceutilsget_device_properties)devicer mem_clock_khz bus_widthbw_gbpss r get_dram_gbpsr^ks  ~3355DDFFM'==fEEFVWM #99&AA/RIi'!+c1A5G Nrcddl}ddlm}|s|j}|jj|ddz}|j|}|ddkr||j ksJd}ni||j |j fvrd}nV||j |j |j fvrd}n=||jtjtjtjfvrd }nt'd ||z|zd z}|S) NrrrSmultiprocessor_countrViidtype not supported& .>)r{r rr|rWrrXrYget_device_capabilityfloat16float32int32rint16int8tl float8e4nv float8e4b15float8e5r r clock_raterZr{r num_subcores capabilityops_per_sub_coretflopss rget_max_tensorcore_tflopsruws(LLL -**,,=&<.decorator..wrappersW MMMrz||4499;;I - 3 3 5 5 G  )Y/%A%Aw''(;J(GHH!z&1UXYY F***,n*** +09<>>!1>>G>>> nox%L]agjkkk~***,e***0C OOCCCCCC((((((r) functoolswraps)rrPrs` r decoratorz cuda_memcheck..decorators>  ! ! ) ) ) ) ) " ! )"rr)rrs` r cuda_memcheckrs$, rFc #K tjgdtjdddd|d|gtjdddd|d|gtdgd }td gd }t||z d ksJd |d t||z d ksJd |d d|z}d|zdz}||fVtjgdtjgdtjgddS#tjgdtjgdtjgdwxYw)N)r r r-pmr{r r rz--lock-gpu-clocks=r z--lock-memory-clocks=zclocks.current.smrzclocks.current.memoryrzzGPU SMs must run at z MHzg3O?igMbP?)r r rrr)r r rz-rgc)r r rz-rmc)rrr#abs) ref_sm_clock ref_mem_clock cur_sm_clock cur_mem_clockrtgbpss r set_gpu_clockrsC E E EFFF    > > > > > !        CM C CM C C !    1233A6 6788; <,.//"4446_\6_6_6_444==011B6668b}8b8b8b666)L8&-dl E E EFFF A A ABBB A A ABBBBB  E E EFFF A A ABBB A A ABBBBs CD!!AE%cddl}ddlm}|s|j}|jj|ddz}|j}|ddkr+||j krd}nM||j krd}n?td ||j krd}n"||j |j fvrd}ntd ||z|zd z}|S) NrrrSr`rarV @rcrd) r{r rr|rWrrXrYrergrfrrros rget_max_simd_tflopsrsLLL -**,,=&<rsx   %%%%%%"""""""""" * * * ((($ ) ) )5550@B@B@B@BFIDIDIDIDXG@G@G@G@TH@H@H@H@V0^0^0^0^f@@@@@@@@F````````F    :6CCCC8r