o Zh+@s,ddlZddlZddlmZmZddlmZddlmZddl m Z m Z ddl m Z mZmZmZddlZddlmZmZddlmZejed ZeoOejZd Zed Zed Zd e e e efefde e e efeffddZ GdddZ!Gddde!Z"Gddde"Z#ere#Z$dSe"Z$dS)N)cached_propertywraps)chain)median)AnyCallable) Concatenate ParamSpecSelfTypeVar)counters dynamo_timed)use_experimental_benchmarkerZ benchmarkingiPTfnreturnc s.tdtdtjdtjdtffdd }|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. selfargskwargsrcst|jjdj}tdd|d7<t|dd|g|Ri|WdS1s3wYdS)N.Zinductorz benchmarking.T)Zlog_pt2_compile_event) __class____name__r r )rrrZ fn_qual_namerS/var/www/auris/lib/python3.10/site-packages/torch/_inductor/runtime/benchmarking.pywrapper"s $ztime_and_count..wrapper)rrrrrr)rrrrrtime_and_counts$rc @seZdZdeddfddZedededefdeedfd e e efd ede f d d Z e ddedegefde de de f ddZededed ede fddZdS) BenchmarkerrrNcCsdSNr)rrrr__init__-szBenchmarker.__init__r.fn_args fn_kwargsrc sd}tD]}t|tjsq |dur|j}q |j|kr#tdq |dur,tdfdd}|tdkrD|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.csiSr rrrr"r#rrYsz'Benchmarker.benchmark..cpu) rvalues isinstancetorchZTensordevice ValueError benchmark_cpu benchmark_gpu)rrr"r#rZinferred_deviceZ arg_or_kwarg _callablerr$r benchmark0s&  zBenchmarker.benchmarkdr.warmuprepcs.dtdttffdd }||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. msrcsNg}t} t}t}|||t||t|kr& |Sqr )time perf_counterappendMILLISECONDS_PER_SECOND)r4ZtimingsZ run_start_tZstart_tZend_tr.rrrun_forusz*Benchmarker.benchmark_cpu..run_for)intlistfloatr)rr.r2r3r:rr9rr,as zBenchmarker.benchmark_cpurcOstr )NotImplementedError)rrrrrrr-szBenchmarker.benchmark_gpu)r0r1)r __module__ __qualname__r r!rrrtupledictstrr=r/r;r,r-rrrrr,s>    0  " rc @sPeZdZedededeffddZededegefdede fdd Z d S) TritonBenchmarkerrr.c Cs6z ddlm}W|Sty}ztd|d}~ww)z"Lazily import Triton's `do_bench`.r)do_benchzrequires TritonN)Ztriton.testingrE ImportErrorr>)rrEerrrtriton_do_benchs z!TritonBenchmarker.triton_do_benchr.rcKst|jj}t|D] }||vr||=q d|vr&|j|fi|dSd|vr3|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. Z quantilesrZ return_moder)inspect signaturerH parametersr<keys)rr.rZdo_bench_paramskwargrrrr-szTritonBenchmarker.benchmark_gpuN) rr?r@rr rrrHrr=r-rrrrrDs (rDc@seZdZededefddZdededeee j j e j j ffddZ dedeee j j e j j fde fd d Ze ddedegefdededededede fddZdS)InductorBenchmarkerrrcCstj}tj|}|jS)z7Get the L2 cache size, in bytes, of the current device.)r)cudaZcurrent_deviceZget_device_properties L2_cache_size)rr*propsrrrrPs  z!InductorBenchmarker.L2_cache_sizeiterscCsddt|DS)z!Get `iters` pairs of CUDA events.cSs(g|]}tjjddtjjddfqS)T)Z enable_timing)r)rOEvent).0_rrr s  z7InductorBenchmarker.get_event_pairs..)range)rrRrrrget_event_pairssz#InductorBenchmarker.get_event_pairs event_pairscCstdd|DS)zIGet the minimum timing, in milliseconds, for a group of CUDA event pairs.cSsg|] \}}||qSr)Z elapsed_time)rT start_event end_eventrrrrVszBInductorBenchmarker.get_event_pairs_min_timing..)min)rrYrrrget_event_pairs_min_timings z.InductorBenchmarker.get_event_pairs_min_timingr1r.estimation_itersmemory_warmup_itersbenchmark_itersmax_benchmark_durationrcKstj|tjtj|jdtjdd}|||}|D]\} } || || q$tj| |} t t |t|| d}t |D]} |qR||}|D]\} } || || q`tj| |} ~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. rO)Zdtyper*r) r)rOZ synchronizeemptyrPr;Zzero_rXrecordr]maxr\rW)rr.r`rarbrcrbufferrYrZr[Zestimated_timingrUZbenchmarked_timingrrrr-s8 !              z!InductorBenchmarker.benchmark_gpuN)r^r1r1r_)rr?r@rr r;rPr<rAr)rOrSrXr=r]rrrr-rrrrrNsN     rN)%rIr5 functoolsrr itertoolsr statisticsrtypingrrZtyping_extensionsrr r r r)Ztorch._dynamo.utilsr r Ztorch._inductor.configrZ_loggingZgetArtifactLoggerrloggerrOZ is_availabler8rrrrrDrNZ benchmarkerrrrrs8     ]'q