Vh+ddlZddlZddlmZmZddlmZddlmZddl m Z m Z ddl m Z mZmZmZddlZddlmZmZddlmZej.j1ed Zexrej6j9Zd Zed Zed Zd e e e efefde e e efeffdZ GddZ!Gdde!Z"Gdde"Z#ere#Z$ye"Z$y)N)cached_propertywraps)chain)median)AnyCallable) Concatenate ParamSpecSelfTypeVar)counters dynamo_timed)use_experimental_benchmarker benchmarkingiPTfnreturnc tdtdtjdtjdt ffd }|S)zWraps `fn` with `dynamo_timed` context, and increments the appropriate dynamo counters. It is expected that `fn` is a method of `Benchmarker` or one of its subclasses; typing limitations prevent us from declaring this directly. selfargskwargsrc|jjdj}tdd|xxdz cc<t|d5|g|i|cdddS#1swYyxYw)N.inductorz benchmarking.T)log_pt2_compile_event) __class____name__r r)rrr fn_qual_namers T/home/dcms/DCMS/lib/python3.12/site-packages/torch/_inductor/runtime/benchmarking.pywrapperztime_and_count..wrapper"so..112!BKK=A }\N;<A< ,d C -d,T,V, - - -s A!!A*)rrrrrr)rr"s` r!time_and_countr#sF 2Y-c-!&&-AHH--- Nc eZdZdeddfdZedededefdeedfde e efd ede f d Z e dded egefd e d e de f dZededed ede fdZy) BenchmarkerrrNcyN)rs r!__init__zBenchmarker.__init__-s r$r.fn_args fn_kwargsrc vd}tjD]F}t|tjs| |j }-|j |k7s=t d| t dfd}|tj dk(r|j|fi|S|j|fi|S)aBenchmark `fn(*fn_args, *fn_kwargs)` and return the runtime, in milliseconds (the actual runtime calculation is dictated by the benchmarking implementation, but may be one of [mean, median, minimum, etc.]). Functions as a convenience wrapper around device-specific implementations, like `benchmark_cpu` and `benchmark_gpu`. Raises `ValueError(...)` if we can't safely infer the device type of `fn`; for example, if multiple device types are found in `fn_args` and `fn_kwargs`, or if no device types are found. Arguments: - fn: The function to benchmark. - fn_args: The function's arguments. - fn_kwargs: The function's kwargs. Keyword Arguments: - **kwargs: The benchmarking implementation's kwargs. Returns: - The runtime of `fn(*fn_args, **fn_kwargs)`, in milliseconds. NzcCan't safely infer the device type of `fn` with multiple device types in `fn_args` and `fn_kwargs`!zCan't safely infer the device type of `fn` with no device types in `fn_args` or `fn_kwargs`! You should be calling `.benchmark_cpu` or `.benchmark_gpu` directly.ciSr(r))rr+r,sr!z'Benchmarker.benchmark..YsB595r$cpu) rvalues isinstancetorchTensordevice ValueError benchmark_cpu benchmark_gpu)rrr+r,rinferred_device arg_or_kwarg _callables ``` r! benchmarkzBenchmarker.benchmark0s6!'9+;+;+=> LlELL9&"."5"5$$7 y    "t 6 ell51 1%4%%i:6: :"t!!)6v66r$r;warmuprepchdtdttffd }||t||S)aBenchmark the CPU callable, `_callable`, and return the median runtime, in milliseconds. Arguments: - _callable: The CPU callable to benchmark. Keyword Arguments: - warmup: Optionally, the duration, in milliseconds, to run `_callable` before benchmarking starts. - rep: Optionally, the duration, in milliseconds, to run `_callable` during benchmarking. Returns: - The median runtime of `_callable`, in milliseconds. msrcg}tj} tj}tj}|j||z tz||z tz|kDr |S]r()time perf_counterappendMILLISECONDS_PER_SECOND)r@timings run_start_tstart_tend_tr;s r!run_forz*Benchmarker.benchmark_cpu..run_foruspG++-K++- ))+3JJK[(,CCrINr$)intlistfloatr)rr;r=r>rJs ` r!r7zBenchmarker.benchmark_cpuas2(  U  gcl##r$rctr()NotImplementedError)rrrs r!r8zBenchmarker.benchmark_gpus!!r$)d)r __module__ __qualname__r r*r#rrtupledictstrrMr<rKr7r8r)r$r!r&r&,s t  .7.7 S#X .7sCx.7S> .7  .7  .7.7`OR $ $'C0 $:= $IL $  $ $D"D"""""r$r&c \eZdZedededeffdZededegefdede fdZ y) TritonBenchmarkerrr.cN ddlm}|S#t$r}td|d}~wwxYw)z"Lazily import Triton's `do_bench`.r)do_benchzrequires TritonN)triton.testingrZ ImportErrorrO)rrZes r!triton_do_benchz!TritonBenchmarker.triton_do_benchs4 @ / @%&78a ? @s $ $r;rc 0tj|jj}t |j D] }||vs||= d|vr|j|fi|dSd|vr|j|fi|S|j|fi|ddiS)aBenchmark the GPU callable, `_callable`, and return the runtime, in milliseconds. Arguments: - _callable: The GPU callable to benchmark. Keyword Arguments: - quantiles: Optionally, a tuple of floats denoting the requested quantiles. - return_mode: Optionally, the requested return mode. Currently, Triton's `do_bench` supports min, max, mean, and median return modes. - **kwargs: Additional kwargs passed to Triton's `do_bench`. Returns: - The runtime of `callable`, in milliseconds. If `kwargs["quantiles"]` is specified, this is the first requested quantile. Else, if `kwargs["return_mode"]` is specified, this is the requested return mode. Otherwise, this is the median. quantilesr return_moder)inspect signaturer^ parametersrLkeys)rr;rdo_bench_paramskwargs r!r8zTritonBenchmarker.benchmark_gpus$"++D,@,@ALL&++-( "EO+5M " & '4'' O#ORWOOr$rXcTeZdZededefdZdededeee jje jjffdZ dedeee jje jjfde fdZe ddedegefd ed ed ed ed ede fdZy)InductorBenchmarkerrrctjj}tjj|}|jS)z7Get the L2 cache size, in bytes, of the current device.)r3cudacurrent_deviceget_device_properties L2_cache_size)rr5propss r!rnz!InductorBenchmarker.L2_cache_sizes6**, 008"""r$itersct|Dcgc]B}tjjdtjjdfDc}Scc}w)z!Get `iters` pairs of CUDA events.T) enable_timing)ranger3rkEvent)rrp_s r!get_event_pairsz#InductorBenchmarker.get_event_pairssV5\     t 4   t 4    sAA event_pairsc ht|Dcgc]\}}|j|c}}Scc}}w)zIGet the minimum timing, in milliseconds, for a group of CUDA event pairs.)min elapsed_time)rrw start_event end_events r!get_event_pairs_min_timingz.InductorBenchmarker.get_event_pairs_min_timings=/: *K((3    s. r;estimation_itersmemory_warmup_itersbenchmark_itersmax_benchmark_durationrc tjj|tjjtj|jdztj d}|j |j|}|D]<\} } |j | j|| j>tjj|j|} tt|t || zd}t|D]} |j |j|}|D]<\} } |j | j|| j>tjj|j|} ~t| | S)a<Benchmark a GPU callable using a custom benchmarking implementation. Arguments: - _callable: The callable to benchmark. Keyword Arguments: - estimation_iters: Optionally, the number of iterations to run `_callable` during runtime estimation. - memory_warmup_iters: Optionally, the number of iterations to flush the L2 cache before starting benchmarking. - benchmark_iters: Optionally, the number of iterations to run `_callable` during the benchmarking. - max_benchmark_duration: Optionally, the maximum duration of the benchmarking, in milliseconds. An estimated duration is calculated based on the values of `memory_warmup_iters` and `benchmark_iters`, along with the estimated runtime of `_callable` and various other factors, and we then shrink `benchmark_iters` to fit in the alloted maximum duration. - **kwargs: Additional kwargs that may be passed to the fallback. Returns: - The minimum runtime of `_callable`, in milliseconds. rk)dtyper5r) r3rk synchronizeemptyrnrKzero_rvrecordr}maxryrs)rr;r~rrrrbufferrwr{r|estimated_timingrubenchmarked_timings r!r8z!InductorBenchmarker.benchmark_gpusB       T//14EIIfU **+;< &1  "K LLN    K         ::;G %;?O%O!P QST  *+ A LLN **?; &1  "K LLN    K         !<<[I #%788r$N)rQrQ)rrRrSrr rKrnrLrTr3rkrtrvrMr}r#rrr8r)r$r!riris##D#S##      eEJJ$$ejj&6&667 8      !%eEJJ,<,N>N,N&O!P    !"#&"&( M9M9BG$M9M9! M9  M9 !$ M9M9 M9M9r$ri)%rbrB functoolsrr itertoolsr statisticsrtypingrrtyping_extensionsr r r r r3torch._dynamo.utilsr rtorch._inductor.configr_logginggetArtifactLoggerrloggerrk is_availablerErrr#r&rXri benchmarkerr)r$r!rs , CC 6?  ) )(N C >UZZ%<%<%>  cN CLS!V$a'( k#q&!1$%$Z"Z"z$O $ONm9+m9b: ?P?R r$