a Р║ФhPEу@s╠UdZddlZddlZddlmZddlmZddlmZm Z m Z mZddlZe e e ge fed<ejабr|ddlmZndZiZeeefed<iZeee fed <Gd dДdГZGdd Дd ГZGddДdeГZe e e ge fed<ejабРrddlmZndZGddДdeГZeGddДdГГZGddДdeГZ GddДdeГZ!iZ"eee#efed<da$eeej%fe#edЬddДZ&eeej%fe#edЬd d!ДZ'ee(ee#efd"Ьd#d$ДZ)d%d&ДZ*dS)'aю Device abstraction layer for TorchDynamo and Inductor backends. This module provides a unified interface for different hardware backends (CUDA, XPU, CPU, MPS) through a common device interface. Key components include: - DeviceInterface: Base class defining the common API for all device types - Device-specific implementations: CudaInterface, XpuInterface, CpuInterface, MpsInterface - Device registration system for managing available backends - Worker APIs for multi-processing scenarios - Stream and event management across different devices - Device property caching for worker processes The abstraction layer enables device-agnostic code in TorchDynamo while allowing specialized implementations for each hardware backend's unique features. щN)┌Iterable)┌ dataclass)┌Any┌Callable┌Optional┌Union┌get_cuda_stream)┌_cuda_getCurrentRawStream┌ caching_worker_device_properties┌caching_worker_current_devicesc@s╥eZdZdZGddДdГZGddДdГZGddДdГZGdd Дd ГZed dДГZ ee jjdЬd dДГZ eeedЬddДГZeeedЬddДГZeddДГZeedЬddДГZee jdЬddДГZeddДГZee jdЬddДГZeeeed Ьd!d"ДГZeeed#Ьd$d%ДГZed:e jjdЬd'd(ДГZed;e jjdЬd)d*ДГZede jeed1Ьd2d3ДГZed?e jjedЬd4d5ДГZ ed@e jjedЬd6d7ДГZ!edAe jjd&dЬd8d9ДГZ"d&S)B┌DeviceInterfacezЮ This is a simple device runtime interface for Inductor. It enables custom backends to be integrated with Inductor in a device-agnostic semantic. c@seZdZejjdЬddДZdS)zDeviceInterface.deviceй┌devicecCstВdSйNй┌NotImplementedErrorй┌clsrйr·L/var/www/auris/lib/python3.9/site-packages/torch/_dynamo/device_interface.py┌__new__/szDeviceInterface.device.__new__N)┌__name__┌ __module__┌__qualname__┌torch┌types┌Devicerrrrrr.src@seZdZddДZdS)zDeviceInterface.EventcOstdГВdS)NzYEvent should be inherited from torch.Event, otherwise, it couldn't be captured by dynamo.rйr┌args┌kwargsrrrr3s zDeviceInterface.Event.__new__Nйrrrrrrrr┌Event2sr!c@seZdZddДZdS)zDeviceInterface.StreamcOstdГВdS)Nz[Stream should be inherited from torch.Stream, otherwise, it couldn't be captured by dynamo.rrrrrr9s zDeviceInterface.Stream.__new__Nr rrrr┌Stream8sr"c@sLeZdZdZeedЬddДГZeedЬddДГZedej j dЬd d ДГZdS)zDeviceInterface.Workera Worker API to query device properties that will work in multi processing workers that cannot use the GPU APIs (due to processing fork() and initialization time issues). Properties are recorded in the main process before we fork the workers. r cCstВdSrrr rrr┌ set_deviceFsz!DeviceInterface.Worker.set_deviceй┌returncCstВdSrrrrrr┌current_deviceJsz%DeviceInterface.Worker.current_deviceNcCstВdSrrr rrr┌get_device_propertiesNsz,DeviceInterface.Worker.get_device_properties)N)rrr┌__doc__┌staticmethod┌intr#r&rrrr'rrrr┌Worker>sr+cCstВdSrrrrrrr&RszDeviceInterface.current_devicer cCstВdSrrr rrrr#VszDeviceInterface.set_deviceйrr%cCstВdSrrr rrr┌maybe_exchange_deviceZsz%DeviceInterface.maybe_exchange_devicecCstВdSrrr rrr┌exchange_device^szDeviceInterface.exchange_devicecCstВdSrrrrrr┌device_countbszDeviceInterface.device_countr$cCstВdSrrrrrr┌is_availablefszDeviceInterface.is_availableй┌streamcCstВdSrrr1rrrr2jszDeviceInterface.streamcCstВdSrrrrrr┌current_streamnszDeviceInterface.current_streamcCstВdSrrr1rrr┌ set_streamrszDeviceInterface.set_streamйZ stream_idZdevice_indexZdevice_typecCstВdSrrr5rrr┌_set_stream_by_idvsz!DeviceInterface._set_stream_by_id)┌ device_idxr%cCstВdSrrйr7rrr┌get_raw_streamzszDeviceInterface.get_raw_streamNcCstВdSrrr rrr┌synchronize~szDeviceInterface.synchronizecCs|jа|бSr)r+r'rrrrr'Вsz%DeviceInterface.get_device_propertiescCstВdSrrr rrr┌get_compute_capabilityЖsz&DeviceInterface.get_compute_capabilityFй┌including_emulationcCstВdSrrr<rrr┌is_bf16_supportedКsz!DeviceInterface.is_bf16_supportedй┌dtyper=r%cCs|tjkp|а|бSr)r┌bfloat16r>йrr@r=rrr┌is_dtype_supportedОsz"DeviceInterface.is_dtype_supportedcCstВdSrrr rrr┌memory_allocatedФsz DeviceInterface.memory_allocatedcCsdS)zТ Returns True if the device has Triton support, False otherwise, even if the appropriate Triton backend is not available. Frr rrr┌is_triton_capableШsz!DeviceInterface.is_triton_capablecCs|абstdГВdS)aH Raises a `RuntimeError` with the appropriate human-readable instructions to resolve the issue if Triton is not available for the given device, or the default device if `device` is `None`. The caller should ensure the presence of the 'triton' package before calling this method. z/This device is not capable of supporting TritonN)rE┌RuntimeErrorrrrr┌raise_if_triton_unavailableаs z+DeviceInterface.raise_if_triton_unavailable)N)N)N)F)F)N)N)N)#rrrr(rr!r"r+r)r&rrrr#r*r-r.r/┌boolr0r2r3r4r6r9r:┌classmethodr'r;r>r@rCrDrErGrrrrr(s^ ■rc@sDeZdZdZeeeeddЬddДZddДZ e e e dЬd d ДZdS)┌DeviceGuarda_ This class provides a context manager for device switching. This is a stripped down version of torch.{device_name}.device. The context manager changes the current device to the given device index on entering the context and restores the original device on exiting. The device is switched using the provided device interface. N)┌device_interface┌indexr%cCs||_||_d|_dS)Nщ )rK┌idx┌prev_idx)┌selfrKrLrrr┌__init__╕szDeviceGuard.__init__cCs|jdur|jа|jб|_dSr)rNrKr.rO)rPrrr┌ __enter__┐s zDeviceGuard.__enter__)┌type┌value┌ tracebackcCs|jdur|jа|jб|_dSйNF)rNrKr-rO)rPrSrTrUrrr┌__exit__├s zDeviceGuard.__exit__)rrrr(rSrrr*rQrRrrWrrrrrJоs ■rJc@s4eZdZejjZejjZejjZGddДdГZe ejj ГZ e ejjГZe ejjГZe ejj ГZ e ejjГZe ejjГZe ejjГZe ejjГZe ejjГZe eГZe ejjГZe ejjГZe ejjГZe ejjГZe edЬddДГZe dejjdЬdd ДГZe dejjed ЬddДГZ e dejjdd Ьd dДГZ!dS)┌ CudaInterfacec@sHeZdZeedЬddДГZeedЬddДГZed ejj dЬdd ДГZ dS)zCudaInterface.Workerr cCs|td<dSйN┌cudaйrr rrrr#╥szCudaInterface.Worker.set_devicer$cCsdtvrtdStjабSrY)rrrZr&rrrrr&╓sz#CudaInterface.Worker.current_deviceNcCsВ|durшs z>CudaInterface.Worker.get_device_properties..) ┌ isinstance┌strrrrSrLrXr+r&r ┌rangerZr/йrZdevice_proprrrr'▄s ■z*CudaInterface.Worker.get_device_properties)Nйrrrr)r*r#r&rrrr'rrrrr+╤sr+r$cCs tjабSr)rrZr0rrrrr0szCudaInterface.is_availableNr cCsFtjjdur(tjа|б\}}|d|Stjа|бjаddбdSdS)Nщ ·:щr)r┌version┌hiprZ┌get_device_capabilityr'ZgcnArchName┌split)r┌major┌minrrrr;sz$CudaInterface.get_compute_capabilityr,cCstjjduptjа|бjdkS)Nщ)rrhrirZr'rlr rrrrEs■zCudaInterface.is_triton_capablecCsrddlm}tа|бs0tjа|б}||tабГВddl }tj jdurZd|jjvrnt dГВnd|jjvrnt dГВdS)Nr)┌GPUTooOldForTritonZamdz'triton not built with the 'amd' backendZnvidiaz*triton not built with the 'nvidia' backend)Ztorch._inductor.excrorXrErrZr'┌inspect┌currentframe┌triton.backendsrhri┌backendsrF)rroZdevice_props┌tritonrrrrGs z)CudaInterface.raise_if_triton_unavailable)N)N)N)"rrrrrZrr!r"r+r)r&r#r/r2r3r4r6r:r'rr9┌_exchange_devicer.┌_maybe_exchange_devicer-rDr>rHr0rrr;rErGrrrrrX╔s4rX┌get_xpu_stream)┌_xpu_getCurrentRawStreamc@s>eZdZejjZejjZejjZGddДdГZe ejj ГZ e ejjГZe ejjГZe ejj ГZ e ejjГZe ejjГZe ejjГZe ejjГZe ejjГZe eГZe ejjГZe ejjГZe ejjГZe edЬddДГZe dejjdЬdd ДГZe deedЬdd ДГZe dejjedЬddДГZ e dejjddЬddДГZ!dS)┌XpuInterfacec@sHeZdZeedЬddДГZeedЬddДГZed ejj dЬdd ДГZ dS)zXpuInterface.Workerr cCs|td<dSйN┌xpur[r rrrr#1szXpuInterface.Worker.set_devicer$cCsdtvrtdStjабSrz)rrr{r&rrrrr&5sz"XpuInterface.Worker.current_deviceNcCsВ|dur.) r`rarrrSrLryr+r&r rbr{r/rcrrrr';s ■z)XpuInterface.Worker.get_device_properties)Nrdrrrrr+0sr+r$cCs tjабSr)rr{r0rrrrr0^szXpuInterface.is_availableNr cCstjа|б}|Sr)rr{rj)r┌ccrrrr;bsz#XpuInterface.get_compute_capabilityFйr=r%cCs tjабSr)rr{r>r<rrrr>gszXpuInterface.is_bf16_supportedr,cCsdSйNTrr rrrrEkszXpuInterface.is_triton_capable)┌evicer%cCs ddl}d|jjvrtdГВdS)Nr┌intelz)triton not built with the 'intel' backendйrrrsrF)rrtrrrrGosz(XpuInterface.raise_if_triton_unavailable)N)F)N)N)"rrrrr{rr!r"r+r)r&r#r/r2r3r4r6r:r'rwr9rur.rvr-rDrHr0rrr;r>rErGrrrrry+s6ryc@seZdZUeed<dS)┌CpuDevicePropertiesZmulti_processor_countN)rrrr*┌__annotations__rrrrrВws rВc@s╓eZdZGddДdejГZGddДdГZeedЬddДГZeded Ьd dДГZ edej jed ЬddДГZ eedЬddДГZeddДГZedej jdЬddДГZedej jed ЬddДГZedej jdd ЬddДГZdS) ┌CpuInterfacec@s.eZdZd ddДZedЬddДZddd ДZdS)zCpuInterface.EventTcCs d|_dS)Ngй┌time)rPZ enable_timingrrrrQ~szCpuInterface.Event.__init__r$cCs|j|jdS)NiшrЕ)rPZ end_eventrrr┌elapsed_timeБszCpuInterface.Event.elapsed_timeNcCstаб|_dSr)rЖ┌perf_counter)rPr2rrr┌recordДszCpuInterface.Event.record)T)N)rrrrQ┌floatrЗrЙrrrrr!}s r!c@s$eZdZedejjdЬddДГZdS)zCpuInterface.WorkerNr cCsddl}|аб}t|ГSйNr)┌multiprocessing┌ cpu_countrВ)rrМrНrrrr'Иsz)CpuInterface.Worker.get_device_properties)N)rrrr)rrrr'rrrrr+Зsr+r$cCsdSr~rrrrrr0ПszCpuInterface.is_availableFr<cCsdSr~rr<rrrr>УszCpuInterface.is_bf16_supportedNr,cCsdSйN┌rr rrrr;Чsz#CpuInterface.get_compute_capabilitycCsdSrЛrr8rrrr9ЫszCpuInterface.get_raw_streamcCsdSrЛrrrrrr&ЯszCpuInterface.current_devicer cCsdSrrr rrrr:гszCpuInterface.synchronizecCsdSr~rr rrrrEзszCpuInterface.is_triton_capablecCs ddl}d|jjvrtdГВdS)Nr┌cpuz'triton not built with the 'cpu' backendrБ)rrtrrrrGлsz(CpuInterface.raise_if_triton_unavailable)F)N)N)N)N)rrrrr!r+r)rHr0r>rrrar;r*r9r&r:rErGrrrrrД|s$ rДc@sЪeZdZedeedЬddДГZedejeedЬddДГZ eedЬd d ДГZ eddДГZedejj edЬddДГZedejj dЬddДГZGddДdГZd S)┌MpsInterfaceFr}cCstjjаddбS)Nщr)rrs┌mpsZis_macos_or_newerr<rrrr>┤szMpsInterface.is_bf16_supportedr?cCs(|tjtjfvrdS|tjkp&|а|бSrV)r┌float64Z complex128rAr>rBrrrrC╕szMpsInterface.is_dtype_supportedr$cCstjjабSr)rrsrУr0rrrrr0└szMpsInterface.is_availablecCsdSrЛrrrrrr&─szMpsInterface.current_deviceNr,cCsdSrОrr rrrr;╚sz#MpsInterface.get_compute_capabilityr cCstjабdSr)rrУr:r rrrr:╠szMpsInterface.synchronizec@s0eZdZedejjdЬddДГZeddДГZdS)zMpsInterface.WorkerNr cCsiSrrr rrrr'╤sz)MpsInterface.Worker.get_device_propertiescCsdSrЛrrrrrr&╒sz"MpsInterface.Worker.current_device)N) rrrr)rrrr'r&rrrrr+╨sr+)F)F)N)N)rrrr)rHr>rIrr@rCr0r&rrrar;r:r+rrrrrС│s" ■ rС┌device_interfacesFйrrKcCst|tjГr|j}|t|<dSr)r`rrrSrХrЦrrr┌register_interface_for_device▐srЧr,cCs>t|tjГr|j}tstГ|tvr,t|Std|ЫЭГВdS)NzNo interface for device )r`rrrS┌_device_initialized┌init_device_regrХrr rrr┌get_interface_for_deviceцsrЪr$cCsts tГtабSr)rШrЩrХ┌itemsrrrr┌ get_registered_device_interfacesЁsrЬcCsxtdtГttjабГD]}td|ЫЭtГqtdtГttjабГD]}td|ЫЭtГqFtdtГtdt Гda dS)NrZzcuda:r{zxpu:rРrУT)rЧrXrbrrZr/ryr{rДrСrШ)r^rrrrЩЎs rЩ)+r(rprЖ┌collections.abcrZdataclassesr┌typingrrrrrr*rГrZZ_is_compiledZtorch._Cr rr ┌dictrarrrJrXr{rxrwryrВrДrСrХrSrШrrЧrЪ┌tuplerЬrЩrrrr┌s@ [L7'