DŽgbdZddlmZddlmZmZmZddlZddlm Z ddlm Z gd Z ee e edfZd efd Zd efd Zd efd Zd efdZd efdZd efdZd efdZd ded dfdZGddZGddZeZead ded efdZGddeZded efdZd efdZ ded dfdZ!d e fdZ"y)!zj This package implements abstractions found in ``torch.cuda`` to facilitate writing device-agnostic code. )AbstractContextManager)AnyOptionalUnionNdevice)amp) is_available synchronizecurrent_devicecurrent_streamstream set_device device_countStream StreamContextEventreturncRtjjjS)z/Returns a bool indicating if CPU supports AVX2.)torch_C_cpu_is_avx2_supportedZ/home/mcse/projects/flask_80/flask-venv/lib/python3.12/site-packages/torch/cpu/__init__.pyrr s 88== + + --rcRtjjjS)z1Returns a bool indicating if CPU supports AVX512.)rrr_is_avx512_supportedrrrr r %s 88== - - //rcRtjjjS)z6Returns a bool indicating if CPU supports AVX512_BF16.)rrr_is_avx512_bf16_supportedrrrr"r"*s 88== 2 2 44rcRtjjjS)z/Returns a bool indicating if CPU supports VNNI.)rrr_is_avx512_vnni_supportedrrr_is_vnni_supportedr%/s 88== 2 2 44rcRtjjjS)z3Returns a bool indicating if CPU supports AMX_TILE.)rrr_is_amx_tile_supportedrrrr'r'5s 88== / / 11rcRtjjjS)zInitializes AMX instructions.)rrr _init_amxrrrr)r):s 88== " " $$rcy)zReturns a bool indicating if CPU is currently available. N.B. This function only exists to facilitate device-agnostic code Trrrrr r ?s rr cy)zWaits for all kernels in all streams on the CPU device to complete. Args: device (torch.device or int, optional): ignored, there's only one CPU device. N.B. This function only exists to facilitate device-agnostic code. Nrrs rr r Hrc*eZdZdZddeddfdZddZy) rzH N.B. This class only exists to facilitate device-agnostic code priorityrNcyNr)selfr.s r__init__zStream.__init__W rcyr0rr1rs r wait_streamzStream.wait_streamZr3r)rN)__name__ __module__ __qualname____doc__intr2r6rrrrrRs  d  rrc4eZdZdefdZdddZddZdddZy) rrcy)NTrr1s rqueryz Event.query_srNcyr0rr5s rrecordz Event.recordbr3rcyr0rr@s rr zEvent.synchronizeer3rcyr0rr5s rwaitz Event.waithr3rr0r8)r9r:r;boolrArCr rFrrrrr^st   rrctS)zReturns the currently selected :class:`Stream` for a given device. Args: device (torch.device or int, optional): Ignored. N.B. This function only exists to facilitate device-agnostic code )_current_streamrs rrrps  rcFeZdZUdZeeed<dZdZde de de dd fd Z y ) rzvContext-manager that selects a given stream. N.B. This class only exists to facilitate device-agnostic code cur_streamc(||_t|_yr0)r_default_cpu_stream prev_streamr5s rr2zStreamContext.__init__s .rc<|j}|yt|_|ayr0)rrIrN)r1rKs r __enter__zStreamContext.__enter__s#[[   +$rtypevalue tracebackrNc:|j}|y|jayr0)rrNrI)r1rQrRrSrKs r__exit__zStreamContext.__exit__s![[   **r) r9r:r;r<rr__annotations__r2rPrrUrrrrr|s?   /%+S++++rrrct|S)zWrapper around the Context-manager StreamContext that selects a given stream. N.B. This function only exists to facilitate device-agnostic code )r)rs rrrs   rcy)zReturns number of CPU devices (not cores). Always 1. N.B. This function only exists to facilitate device-agnostic code r rrrrrrs rcy)zzSets the current device, in CPU we do nothing. N.B. This function only exists to facilitate device-agnostic code Nrrs rrrr,rcy)zyReturns current device for cpu. Always 'cpu'. N.B. This function only exists to facilitate device-agnostic code cpurrrrrrs rr0)#r< contextlibrtypingrrrrr _devicer __all__strr= _device_trGrr r"r%r'r)r r rrrMrIrrrrrrrrrrcs# .''   '3T) * .D. 0d0 545 5D5 22 %4% d T      h% 9  +*+>!6!4! c yTr