fG d dl Z d dlZd dlZd dlZd dlmZ d dlmZmZm Z ddl
m Z d Z
ddZddZdd Z G d
d Z G d d
Zd ZddZddZd Zedd ZddZy) N)contextmanager)AnyDictList )languagec dj | } dddd| z dg}t j | }|j t j
j j d }|D cg c]
}t | }}|S c c}w )N,
nvidia-smi-i0z --query-gpu=z --format=csv,noheader,nounits) join
subprocess check_outputdecodesysstdoutencodingsplitint)attrscmdoutretxs W/var/lib/jenkins/workspace/mettalog/venv/lib/python3.12/site-packages/triton/testing.pynvsmir
sy HHUOE sNU$:<[
\C
!
!#
&C
**SZZ((
)
/
/
4C
a3q6
C
J s -Bc ddl }|dv sJ |j j |j j k( r t d | |/|D ]* }|j |j
d d|_ , |j j }|j j | 5 | ddd |j j |j j d }|j j d }|j |j |j |j j |j | } t dt! || z }
|j j }|j j | 5 t# |
D ] } ||D ] }d|_ | ddd |j j g } d}
t# |
D ] } |j j d }|j j d }|j |j |j |j j | |j | |
z gz
} |j% | } t' || | j) S # 1 sw Y .xY w# 1 sw Y xY w) 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
r NminmaxmeanmedianzQCannot capture graph in default stream. Please use side stream in benchmark code.T
enable_timingr
)torchcudacurrent_streamdefault_stream RuntimeErrordetach_requires_grad_grad CUDAGraphgraph synchronizeEventrecordreplay elapsed_timer! r rangetensorgetattritem)fnrep grad_to_none return_moder' r g start_event end_event estimate_msn_repeatir n_retriestimess r do_bench_cudagraphrF s : :: : zz "ejj&?&?&AAnooD A
IIK
T "AF
A !
JJ **"""6K
t 4IHHJ
JJ **95K1c# +,-H
A ! x A'%A !AF & D !
JJ
CI
9
jj&&T&: JJ$$4$8
((3h>??
LL E &75+ &u - 2 2 44C
s 8K,(K9,K69Lc . |dv sJ ddl } | |j j |r(|j t d |j d }n'|j t d |j
d }|j j
d } |j j
d }
| j t d
D ] } |j | |
j |j j | j |
d
z } t d t || z }
t d t || z }t | D cg c] }|j j
d } }t | D cg c] }|j j
d }
}t |
D ] } | t | D ]O }||D ] }d|_ |j | | j | |
| j Q |j j |j t | |
D cg c] \ }}|j | c}}|j }|P|j! ||j ||j j# }t% | d k( r|d }|S t' || | j) S c c}w c c}w c c}}w )
a
Benchmark the runtime of the provided function. By default, return the median runtime of :code:`fn` along with
the 20-th and 80-th performance percentile.
: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 to None
:type grad_to_none: torch.tensor, optional
:param quantiles: Performance percentile to return in addition to the median.
:type quantiles: list[float]
:param fast_flush: Use faster kernel to flush L2 between measurements
:type fast_flush: bool
r r Ng Ar( )dtypedeviceg ATr$ r )rH )r' r( r1 emptyr int8r2 r3 r6 zero_r5 r! r. r7 zipfloatquantiletolistlenr8 r9 )r: warmupr; r< quantiles
fast_flushr= r' cacher? r@ _rA n_warmuprB rC r serE r s r do_benchr[ R s $ : :: :D JJ
C
O599V L CJejj H **"""6K
t 4I
1X
JJ **959K 1c&;./0H1c# +,-HAFx QA5::##$#7QKQ?DXO!!!!5OIO
8_
8_ #! "
A
!
JJ LLK8ST1!..+T\a\g\gL hE nnUELL%++L$NOVVX s8q=a&C
&75+ &u - 2 2 447 RO( Us '#L#L &L
c dd l }dd l}t | |j s|j | } t ||j s|j | }|d}t | r || j n|}|d}t | r || j n|}t | |j rU| j |j k( r| j } | j j j } t ||j rU|j |j k( r|j }|j j j }| j dkD s|j dkD r!|j j | |||d y |j | ||| st | d| d | d
| d | d
y )
Nr g{Gz?g r T)atolrtol equal_nan)r] r^ z is not close to z (atol=z, rtol=))numpyr'
isinstanceTensorr7 callablerH bfloat16rO cpudetachsizetestingassert_allcloseallcloseAssertionError)r yr] r^ err_msgnpr' s r assert_closerq s a & LLO a & LLO | $TN4=D | $TN4=D !U\\" 77enn $ A
EEGNN " " $!U\\" 77enn $ A
EEGNN " " $ vvzQVVaZ
""1ad"N
;;q!$T; 2 y!,=aSvWUYTZZ[ \]] 3 c t e Zd ZdZ ddee dee dedee dee ded eeef d
ed ed ed
efdZ y) Benchmarkzk
This class is used by the :code:`perf_report` function to generate line plots with a concise API.
Nx_namesx_valsline_arg line_vals
line_names plot_nameargsxlabelylabelx_logy_logc || _ || _ |
| _ || _ || _ || _ | | _ |
| _ || _ | | _ || _
|| _ y)a
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
N) ru rv r~ rw rx ry r stylesr| r} rz r{ )selfru rv rw rx ry rz r{ r| r} r~ r colorr s r __init__zBenchmark.__init__ sY \
"$
" rr ) r FFNN)
__name__
__module__ __qualname____doc__r strr r boolr rr r rt rt s :c: S : :
9 : I
: : 38n: : : : :rr rt c 8 e Zd Zd Z d dedededefdZd
dZy) Markc || _ || _ y Nr:
benchmarks)r r: r s r r z
Mark.__init__ s $rr bench save_path
show_plots
print_datac dd l }dd lm} dd l}
|j } |j D cg c] } | d }
} |j D cg c] } | d }} t |j }|
j || z |
z |z }|j D ] } t | t
t f s|D cg c] }| } }t | t | k7 rt dt | d| t t || }g g g }}}|j D ]I } | j di ||j" |i|j$ |} |\ } }
}|| gz
}||
gz
}||gz
}K t | |z |z |z |j( t | < |j* r| j- | j/ }|d }t1 |j D ] \ }}||dz ||dz }}
|j2 r|j2 | d nd }|j2 r|j2 | d nd } |j5 || || ||| |
j7 j9 r|j7 j9 r|
j; t< }
|j; t< }|j? || |
|d |
|jA |jC |jD xs | |jG |jH |jK |jL rd nd |jO |jP rd nd |r| jS |r8| jU |jV jY ||j* d
|||j z }|r=|jZ d dk( r+|j\ j_ \ } } || || z
|d< |r1ta |j* dz ta |jc |r?|je |jV jY ||j* d d| dd |S c c} w c c} w c c}w # t& $ r
|d d }}
} Y 8w xY w)Nr z-minz-max)columnsz Expected z
values, got r )labelr lsg333333?)alphar loglinearz.png Diff:z.csvz%.fF) float_formatindexr )3osmatplotlib.pyplotpyplotpandasry listru DataFramerv rc tuplerR
ValueErrordictrN rx r: rw r{ TypeErrorlocrz figuresubplot enumerater plotisnullallastyperO fill_betweenlegend
set_xlabelr|
set_ylabelr}
set_xscaler~
set_yscaler showsavefigpathr shaper rQ print to_stringto_csv) r r r r r diff_colsave_precisionkwragsr pltpdy_meanr y_miny_maxru dfrW x_argsrow_meanrow_minrow_maxrn r axfirst_xrC colstycol0col1s r _runz Mark._run s '!!%*%5%56A3d66%*%5%56A3d66u}}%
\\'F"2U":U"B\
C A a$ / '(1Q((1vW % 9S\N-s!KLL#gq/*F)+RwgH__ dggVV5>>1*=VVvV;+.(FE5 VH$E7"E7" % #1g 07 :W DBFF3r7O' * ?? JJL B ajG!%"2"231 !!f*~r!f*~u,1LLell1oa(d,1LLell1oa(d7 RU!33G ||~))+ELLN4F4F4H !LL /E !LL /EOOBwKTQTOU 4
IIK MM%,,1' 2 MM%,, ' MM5;;%H = MM5;;%H =
BGGLLu6Gt4L MN
%***
+ q(**,JD$DBtH,BvJ %//C' ( ",,. ! IIbggll9 0A.FGXZ[iZjjkVl !
# y 76 ) ! ;+.d5EF;s# Q Q#, Q(Q--R ?R c t | j t }|r
| j gn | j }g }|rRt j |d t t j j |d d } | j d |D ]I }
|j | j |
|||fi | |s+ j d|
j d K |r! j d | j |r |r|d S |S y )
NT)exist_okz results.htmlwz
z
z
r )
rc r rt r makedirsopenr r writeappendr rz close) r r r r return_dfkwargshas_single_benchr
result_dfshtmlr s r runzMark.runP s %dooyA*:doo&
KK D 1 Y?ED JJ' ( E idiiy*j [TZ [ \
]5??*;:FG JJ) * JJL !!}$!!rr N)F )FFr F) r r r r rt r r r r r rr r r r s> % ch C) C C CSW CJrr r c fd}|S )z
Mark a function for benchmarking. The benchmark can then be executed by using the :code:`.run` method on the return value.
:param benchmarks: Benchmarking configurations.
:type benchmarks: List of :class:`Benchmark`
c t | S r )r r s r z perf_report..o s b*-rr r )r wrappers ` r perf_reportr h s .G Nrr c ddl }ddlm} | s|j j } |j
j j | d }|j
j j | d }||z dz dz d z }|S )
z return DRAM bandwidth in GB/s r Nr drivermem_clock_rate
mem_bus_widthr g .A )r' runtimer r( current_deviceactiveutilsget_device_properties)rI r' r
mem_clock_khz bus_widthbw_gbpss r
get_dram_gbpsr s sz **,MM''==fEFVWM
##99&A/RIi'!+c1A5G Nrr c J dd l }ddlm} |s|j j }|j
j j | d dz }|j j | }|d dk r| |j k( sJ d}n| |j |j fv rd}nr| |j |j |j fv rd}nJ| |j t j t j" t j$ fv rd }n t' d
||z |z d z }|S ) Nr r r multiprocessor_count r i i dtype not supported& .>)r' r r r( r r r r get_device_capabilityfloat16float32int32rf int16rL tl
float8e4nv float8e4b15float8e5r+ rH
clock_raterI r' r num_subcores
capabilityops_per_sub_coretflopss r get_max_tensorcore_tflopsr s **,==&&<.decorator..wrapper s! rzz|499;I - 3 3 5 G Y/%Aww''(;(;J(GH !zz& 1UXY F*n,nn* +0099<< b!1!1 2!G9A> nnox%L]agjk~~*e,ee*0C
OCCC((rr ) functoolswraps)r- r r, s ` r decoratorz cuda_memcheck..decorator s% ! )
" )" rr r )r, r0 s ` r
cuda_memcheckr1 s , rr c # K t j g d t j dddd| d| g t j dddd| d| g t dg d }t d
g d }t || z
d k s J d | d
t ||z
d k s J d | d
d| z }d|z dz }||f t j g d t j g d t j g d y # t j g d t j g d t j g d w xY ww)N)r r r
-pmr r r r
z--lock-gpu-clocks=r
z--lock-memory-clocks=zclocks.current.smr zclocks.current.memoryr& zGPU SMs must run at z MHzg 3O?i gMbP?)r r r
r3 r
)r r r
z-rgc)r r r
z-rmc)r r r abs) ref_sm_clock
ref_mem_clock cur_sm_clock
cur_mem_clockr gbpss r
set_gpu_clockr: sm C EF a ~ > !
#M?!M? C !
1 23A6 6 78;
<,./"4_8L\NZ^6__4==01B6b:N}o]a8bb6)L8 &-dl EF AB AB EF AB ABs EB>D AEAEEc dd l }ddlm} |s|j j }|j
j j | d dz }|j j }|d dk r/| |j k( rd}nW| |j k( rd}nEt d | |j k( rd}n(| |j |j fv rd}n t d ||z |z d
z }|S ) Nr r r r r r @ r r )
r' r r r( r r r r r r r r+ rf r s r get_max_simd_tflopsr> s **,==&&< r rr r rD s}
% " " <5~I5X"^J? ?D` `F
:6 C C8rr