o m9:jT@s:UdZddlmZddlZddlZddlmZddlmZm Z m Z ddl Z ddl Z ddl mZmZe r?ddlmZddlmZd d lmZd d lmZmZmZmZmZd d lmZmZd ae Z!e"Z#ga$de%d<e&e j'dddZ(ea)dZ*de%d<dqddZ+e+re j'j,Z,e j'j-Z.e j'j/Z0nedZ,drddZ.drdd Z0ed d!dsd"d#Z1dqd$d%Z2dtdud(d)Z3dqd*d+Z4d,d-Z5dvd/d0Z6dvd1d2Z7dvd3d4Z8Gd5d6d6Z9Gd7ddZ:Gd8d9d9e:Z;dwd;d<Zd?Z=eddxdzdAdBZ> dxd{dCdDZ?dsdEdFZ@d|dIdJZAd}dLdMZBGdNdOdOZCd~dRdPZDdvdSdTZEddVdWZFdxddXdYZGdxdd[d\ZHdxdwd]d^ZIdd`daZJddbdcZKddedfZL gdddidjZMdddkdlZNd dmlOmPZPmQZQmRZRmSZSmTZTmUZUmVZVmWZWmXZXmYZYmZZZm[Z[m\Z\m]Z]m^Z^m_Z_m`Z`d dnlambZbmcZcmdZdmeZemfZfmgZgmhZhmiZimjZje doekelelfZmgdpZndS)z This package introduces support for the XPU backend, specifically tailored for Intel GPU optimization. This package is lazily initialized, so you can always import it, and use :func:`is_available()` to determine if your system supports XPU. ) annotationsN) lru_cache)AnyNewType TYPE_CHECKING) _dummy_type_LazySeedTracker)Callable)Device)_get_device_index)graphgraph_pool_handleis_current_stream_capturingmake_graphed_callablesXPUGraph)EventStreamFz*list[tuple[Callable[[], None], list[str]]] _queued_calls_xpu_isInBadForkcCsdSNFrrrY/home/nk/hobo-godmode/plappi-mvp/.venv/lib/python3.10/site-packages/torch/xpu/__init__.py,srrztuple[torch._C.Generator]default_generatorsreturnboolcCstjjS)z(Return true if compile with XPU support.)torch_C_has_xpurrrr _is_compiled1sr _XpuDevicePropertiesdeviceintcCtdNz(PyTorch was compiled without XPU supportNotImplementedErrorr"rrr_exchange_device>r)cCr$r%r&r(rrr_maybe_exchange_deviceAr*r+)maxsizecCstsdStjS)z*Return the number of XPU device available.r)r rr_xpu_getDeviceCountrrrr device_countEs r.cCs tdkS)z7Return a bool indicating if XPU is currently available.r)r.rrrr is_availableMs r/Tincluding_emulationcCstsdS|p tjjS)zKReturn a bool indicating if the current XPU device supports dtype bfloat16.F)r/rxpuget_device_propertieshas_bfloat16_conversions)r0rrris_bf16_supportedSs  r4cCstsdStjjS)zGReturn a bool indicating if the current XPU device supports dtype tf32.F)r/rr1r2'has_subgroup_matrix_multiply_accumulaterrrris_tf32_supported]s r6cCs tot S)z8Return whether PyTorch's XPU state has been initialized.) _initialized_is_in_bad_forkrrrris_initializedhs r9NonecKsftr|dS|ddrt|tdS|ddr(t|tdSt|tfdS)Nseed_allFseed) r9get_lazy_seed_trackerqueue_seed_all traceback format_stack queue_seedrappend)callablekwargsrrr _lazy_callms   rFcCs tdS)zInitialize PyTorch's XPU state. This is a Python API about lazy initialization that avoids initializing XPU until the first time it is accessed. Does nothing if the XPU state is already initialized. N) _lazy_initrrrrinit{s rHc Cststtdr dStqtr WddStr tdts'tdtj dt_ t ddtDz1t D]'\}}z|Wq>tye}zdt|dd |}t||d}~wwWttdnttdwdaWddS1swYdS) Nis_initializingzuCannot re-initialize XPU in forked subprocess. To use XPU with multiprocessing, you must use the 'spawn' start methodz#Torch not compiled with XPU enabledTcss|]}|r|VqdSNr).0callsrrr sz_lazy_init..z5XPU call failed lazily at initialization with error: z' XPU call was originally invoked at: )r9hasattr_tls_initialization_lockr8 RuntimeErrorr AssertionErrorrr _xpu_initrIrextendr> get_calls Exceptionstrjoindelattrr7) queued_callorig_tracebackemsgrrrrGs>      "rGc@s(eZdZdddZddZdd dZdS) _DeviceGuardindexr#rr:cCs||_d|_dSN)idxprev_idx)selfr`rrr__init__s z_DeviceGuard.__init__cCtj|j|_dSrJrr1r)rcrdrerrr __enter__z_DeviceGuard.__enter__typervaluer@cCtj|j|_dSrrr1r+rdrcrerlrmr@rrr__exit__z_DeviceGuard.__exit__N)r`r#rr:rlrrmrr@r)__name__ __module__ __qualname__rfrjrqrrrrr_s r_c@s,eZdZdZdddZddZdd d ZdS)r"zContext-manager that changes the selected device. Args: device (torch.device or int or str): device index to select. It's a no-op if this argument is a negative integer or ``None``. rrr:cCst|dd|_d|_dS)NToptionalrb)r rcrd)rer"rrrrfs zdevice.__init__cCrgrJrhrirrrrjrkzdevice.__enter__rlrmr@cCrnrrorprrrrqrrzdevice.__exit__N)r"rrr:rs)rtrurv__doc__rfrjrqrrrrr"s  cs"eZdZdZdfdd ZZS) device_ofaContext-manager that changes the current device to that of given object. You can use both tensors and storages as arguments. If a given object is not allocated on a XPU, this is a no-op. Args: obj (Tensor or Storage): object allocated on the selected device. rr:cs"|jr|nd}t|dSra)is_xpu get_devicesuperrf)reobjrc __class__rrrfszdevice_of.__init__rr:)rtrurvryrf __classcell__rrrrrzs rzr cCs*tt|}|dkrtj|dSdS)zSet the current device. Args: device (torch.device or int or str): selected device. This function is a no-op if this argument is negative. rN)rGr rr_xpu_setDevicer(rrr set_devices rrXcCs t|jS)aGet the name of a device. Args: device (torch.device or int or str, optional): device for which to return the name. This function is a no-op if this argument is a negative integer. It uses the current device, given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). Returns: str: the name of the device )r2namer(rrrget_device_names rdict[str, Any]cs:t|tttttdtttffddt DS)aGet the xpu capability of a device. Args: device (torch.device or int or str, optional): device for which to return the device capability. This function is a no-op if this argument is a negative integer. It uses the current device, given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). Returns: dict[str, Any]: the xpu capability dictionary of the device Ncs0i|]}|dstt|r|qS)__) startswith isinstancegetattr)rKkeypropsserializable_typesrmrr sz)get_device_capability..) r2r#floatrrXrllisttupledictdirr(rrrget_device_capabilitys rcCstt|dd}t|S)aGet the properties of a device. Returns _XpuDeviceProperties containing the following device properties: - ``name`` (str): device name. - ``platform_name`` (str): SYCL platform name. - ``vendor`` (str): device vendor. - ``device_id`` (int): device identifier (product ID). - ``driver_version`` (str): driver version. - ``version`` (str): runtime version. - ``max_compute_units`` (int): number of parallel compute units. - ``gpu_eu_count`` (int): number of EUs (Execution Unit). - ``max_work_group_size``: (int): maximum number of work-items permitted in a work-group. - ``max_num_sub_groups`` (int): maximum number of sub-groups supported in a work-group. - ``memory_clock_rate`` (int) maximum clock rate of device's global memory in MHz. - ``memory_bus_width`` (int) maximum bus width between device and memory in bits. - ``sub_group_sizes``: (list[int]): a list of supported sub-group sizes. - ``local_mem_size`` (int): device local memory capacity that can be allocated per work-group in bytes. - ``has_fp16`` (bool): whether float16 dtype is supported. - ``has_fp64`` (bool): whether float64 dtype is supported. - ``has_atomic64`` (bool): whether 64-bit atomic operations are supported. - ``has_bfloat16_conversions`` (bool): whether bfloat16 conversions are supported. - ``has_subgroup_matrix_multiply_accumulate`` (bool): whether DPAS (Dot Product Accumulate Systolic) is supported. - ``has_subgroup_matrix_multiply_accumulate_tensor_float32`` (bool): whether DPAS with tf32 inputs is supported. - ``has_subgroup_2d_block_io`` (bool): whether 2D block I/O for efficient matrix multiplication is supported. - ``total_memory`` (int): device global memory in bytes. - ``gpu_subslice_count`` (int): number of subslice. - ``architecture`` (int): device architecture identifier (experimental). - ``type`` (str): device type, e.g. 'cpu', 'gpu', accelerator', 'host', 'unknown'. - ``uuid`` (Any): device UUID (Universal Unique ID), 16 bytes. Args: device (torch.device or int or str): device for which to return the properties of the device. Returns: _XpuDeviceProperties: the properties of the device Trw)rGr _get_device_propertiesr(rrrr2s' r2cCsttjS)z0Return the index of a currently selected device.)rGrr_xpu_getDevicerrrrcurrent_device=s rint | str | torch.device torch.devicecCs2t|tr t|}|St|trtd|}|S)zReturn the torch.device type object from the passed in device. Args: device (torch.device or int or str): selected device. r1)rrXrr"r#r(rrr _get_deviceCs   rpeercCs,tt|dd}t|dd}tj||S)a/Query whether a device can access a peer device's memory. Args: device (torch.device or int or str): selected device. peer (torch.device or int or str): peer device to query access to. Returns: bool: ``True`` if ``device`` can access ``peer``, ``False`` otherwise. Trw)rGr rr_xpu_canDeviceAccessPeer)r"rrrrcan_device_access_peerPs  rc@s6eZdZUdZded<dddZd d ZdddZdS) StreamContextaContext-manager that selects a given stream. All XPU kernels queued within its context will be enqueued on a selected stream. Args: Stream (Stream): selected stream. This manager is a no-op if it's ``None``. .. note:: Streams are per-device. torch.xpu.Stream | None cur_streamstreamrr:cCs*||_tdd|_|jdurd|_dSdS)NTrb)rr rc)rerrrrrfns    zStreamContext.__init__cCs|j}|dus |jdkrdStjd|_|jj|jkr9t|jtj|j|_Wdn1s4wYtj|dSra) rrcrr1current_streamsrc_prev_streamr"dst_prev_stream set_stream)rerrrrrjts zStreamContext.__enter__rlrrmr@cCsJ|j}|dus |jdkrdS|jj|jkrtj|jtj|jdSra)rrcrr"rr1rr)rerlrmr@rrrrrqs zStreamContext.__exit__N)rrrr:rs)rtrurvry__annotations__rfrjrqrrrrr`s   rrrcCst|S)zWrap around the Context-manager StreamContext that selects a given stream. Arguments: stream (Stream): selected stream. This manager is a no-op if it's ``None``. )rrrrrrscCstjj|||ddS)aset stream specified by the stream id, device index and device type Args: stream_id (int): not visible to the user, used to assigned to the specific stream. device_index (int): selected device index. device_type (int): selected device type.  stream_id device_index device_typeN)rr_xpu_setStreamrrrr_set_stream_by_ids  rrcCs*|durdStt|j|j|jddS)aSet the current stream. This is a wrapper API to set the stream. Usage of this function is discouraged in favor of the ``stream`` context manager. Args: stream (Stream): selected stream. This function is a no-op if this argument is ``None``. Nr)rGrrrrrrrrrs  rcCs4ttjt|dd}t|d|d|ddS)aRReturn the currently selected :class:`Stream` for a given device. Args: device (torch.device or int, optional): selected device. Returns the currently selected :class:`Stream` for the current device, given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). Trwrr r)rGrr_xpu_getCurrentStreamr r)r" streamdatarrrrs  rdata_ptrcCs6ttj|t|dd}t|d|d|ddS)a;Return a :class:`Stream` from an external SYCL queue. This function is used to wrap SYCL queue created in other libraries in order to facilitate data exchange and multi-library interactions. .. note:: This function doesn't manage the queue life-cycle, it is the user responsibility to keep the referenced queue alive while this returned stream is being used. The different SYCL queue pointers will result in distinct :class:`Stream` objects, even if the SYCL queues they dereference are equivalent. Args: data_ptr(int): Integer representation of the `sycl::queue*` value passed externally. device(torch.device or int, optional): the device where the queue was originally created. It is the user responsibility to ensure the device is specified correctly. Trwrr rr)rGrr_xpu_getStreamFromExternalr r)rr"rrrrget_stream_from_externals rcCstt|dd}tj|S)a*Wait for all kernels in all streams on a XPU device to complete. Args: device (torch.device or int, optional): device for which to synchronize. It uses the current device, given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). Trw)rGr rr_xpu_synchronizer(rrr synchronizes  r list[str]cCs(tsgStj}|durgS|S)z.)rlenrY) arch_listrrrget_gencode_flagss rtorch._C.GeneratorcCs |j}|dur t}tjj|S)zuReturn the XPU Generator object for the given device. Args: device (torch.device): selected device. N)r`rrr1r)r"rcrrr_get_generators rr1offsetcs$t|dfdd }t|dS)a$Set the random number generator state offset of the specified GPU. Args: offset (int): The desired offset device (torch.device or int, optional): The device to set the RNG state. Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). rr:cst}|dSrJ)r set_offset)default_generator final_devicerrrcbsz!_set_rng_state_offset..cbNr)rrF)rr"rrrr_set_rng_state_offset s  rcCstt|}t|}|S)aLReturn the random number generator state offset of the specified GPU. Args: device (torch.device or int, optional): The device to return the RNG state offset of. Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). .. warning:: This function eagerly initializes XPU. )rGrr get_offset)r"rrrrr_get_rng_state_offsets r)change_current_allocator empty_cacheget_per_process_memory_fractionmax_memory_allocatedmax_memory_reserved mem_get_infomemory_allocatedmemory_reservedmemory_snapshot memory_statsmemory_stats_as_nested_dictMemPoolreset_accumulated_memory_statsreset_peak_memory_statsset_per_process_memory_fraction use_mem_poolXPUPluggableAllocator) get_rng_stateget_rng_state_all initial_seed manual_seedmanual_seed_allr<r; set_rng_stateset_rng_state_all _POOL_HANDLE)9rrrrrrrrrrr"rzr.rrrrr2rrrrrr rrHrr/r4rr9r6rrrrrrrrrrrrrrrr<r;rrrrrrstreamsr)rr)r"r#rr#)rr#)T)r0rrrr)r"r rr:rJ)r"r rrX)r"r rr)r"r rr!)r"rrr)r"r rr rr)rrrr)rrrr:)r"r rr)rr#r"r rr)rr)rrX)r"rrr)r1)rr#r"rrr:)r"rrr#)ory __future__r threadingr@ functoolsrtypingrrrrtorch._C torch._utilsrrcollections.abcr torch.typesr _utilsr graphsr rrrrrrrr7localrPLockrQrrrrr8r>rr r!_xpu_exchangeDevicer)_xpu_maybeExchangeDevicer+r.r/r4r6r9rFrHrGr_r"rzrrrr2rrrrrrrrrrrrrrrmemoryrrrrrrrrrrrrrrrrrrandomrrrrrr<r;rrrr#r__all__rrrrs                  (    ,   +      L,