B \1@sLdZddlmZmZmZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl Z ddl mZddlmZmZmZmZmZmZmZmZddlZddlZddlmZmZddlmZmZdd l m!Z!m"Z"dd l#m$Z$dd l#m%Z%dd l&m'Z'm#Z#m(Z(dd lm)Z)m*Z*ddl+m,Z-e.ej/0ddZ1dZ2ej34dZ5ddZ6Gddde7Z8Gddde7Z9Gddde"Z:ddZ;dZd d!Z?d"d#Z@d$d%ZAe@ZBd&ZCGd'd(d(eDZEeEZFd)d*ZGeGZHGd+d,d,eDZId-d.ZJGd/d0d0eDZKeKZKGd1d2d2eDZLed3d4ZMGd5d6d6eDZNd7d8ZOd9d:ZPd;d<ZQd=d>ZRd?d@ZSdAdBZTdCdDZUGdEdFdFeDZVGdGdHdHeDZWGdIdJdJeDZXGdKdLdLeDZYGdMdNdNeYZZGdOdPdPeZZ[GdQdRdRej\Z]GdSdTdTeDZ^GdUdVdVe^ej\Z_GdWdXdXeDZ`GdYdZdZeDZad[d\ZbGd]d^d^eDZced_d`dadbdcddgZdGdedfdfeDZedgdhZfe'jge'jhe'jie'jje'jkdiZlGdjdkdkeDZmdldmZndndoZodpdqZpdrdsZqdtduZrdvdwZsdxdyZtdd{d|Zud}d~ZvddZwddZxddZyddZzddZ{ddZ|ddZ}dddZ~dddZdddZdddZddZddZejddZdS)a CUDA driver bridge implementation NOTE: The new driver implementation uses a *_PendingDeallocs* that help prevents a crashing the system (particularly OSX) when the CUDA context is corrupted at resource deallocation. The old approach ties resource management directly into the object destructor; thus, at corruption of the CUDA context, subsequent deallocation could further corrupt the CUDA context and causes the system to freeze in some cases. )absolute_importprint_functiondivisionN)product)c_intbyrefc_size_tc_charc_char_p addressofc_void_pc_float) namedtupledeque)utilsmviewbuf)CudaSupportErrorCudaDriverError)API_PROTOTYPES)cu_occupancy_b2d_size)enumsdrvapi_extras)config serialize)longintZNUMBAPRO_VERBOSE_CU_JIT_LOG)rZlinuxcCstt}t|sttj}t t|d}t |t s>tj }| |tjr|ttj}d}|tj|d||n|t|S)Nz;== CUDA [%(relativeCreated)d] %(levelname)5s -- %(message)s)fmt)loggingZ getLogger__name__rZlogger_hasHandlersstrrZCUDA_LOG_LEVELuppergetattr isinstanceintZCRITICALZsetLevelZ StreamHandlersysstderrZ setFormatterZ FormatterZ addHandlerZ NullHandler)ZloggerZlvlZhandlerrr(8lib/python3.7/site-packages/numba/cuda/cudadrv/driver.py _make_logger,s       r*c@s eZdZdS)DeadMemoryErrorN)r __module__ __qualname__r(r(r(r)r+Dsr+c@s eZdZdS) LinkerErrorN)r r,r-r(r(r(r)r.Hsr.cs$eZdZfddZddZZS) CudaAPIErrorcs"||_||_tt|||dS)N)codemsgsuperr/__init__)selfr0r1) __class__r(r)r3MszCudaAPIError.__init__cCsd|j|jfS)Nz[%s] %s)r0r1)r4r(r(r)__str__RszCudaAPIError.__str__)r r,r-r3r6 __classcell__r(r()r5r)r/Ls r/c Cs|tjdd}|dkrttjdkr:tj}dg}dg}n4tjdkrXtj}dg}dg}ntj}d d g}d d g}|dk rytj |}Wn t k rt d |YnXtj |st d||g}n|ddt ||D}g}g}x`|D]X}y ||}WnBtk r@} z"|tj | || Wdd} ~ XYqX|SqWt|r\tnddd|D} t| dS)NZNUMBAPRO_CUDA_DRIVER0Zwin32z\windows\system32z nvcuda.dlldarwinz/usr/local/cuda/libz libcuda.dylibz/usr/libz /usr/lib64z libcuda.soz libcuda.so.1z+NUMBAPRO_CUDA_DRIVER %s is not a valid pathzrNUMBAPRO_CUDA_DRIVER %s is not a valid file path. Note it must be a filepath of the .so/.dll/.dylib or the drivercSsg|]\}}tj||qSr()ospathjoin).0xyr(r(r) yszfind_driver.. css|]}t|VqdS)N)r!)r=er(r(r) szfind_driver..)r:environget_raise_driver_not_foundr&platformctypesZWinDLLZCDLLr;abspath ValueErrorisfilerOSErrorappendallr<_raise_driver_error) ZenvpathZdlloaderZdldirZdlnamesZ candidatesZpath_not_existZdriver_load_errorr;ZdllrBerrmsgr(r(r) find_driverVsN      rQz CUDA driver library cannot be found. If you are sure that a CUDA driver is installed, try setting environment variable NUMBAPRO_CUDA_DRIVER with the file path of the CUDA driver shared library. zM Possible CUDA driver libraries are found but error occurred during load: %s cCs ttdS)N)rDRIVER_NOT_FOUND_MSGr(r(r(r)rFsrFcCstt|dS)N)rDRIVER_LOAD_ERROR_MSG)rBr(r(r)rOsrOcCs>d}t}x,ttD] }||rtt|}|||<qW|S)NZ CUDA_ERROR)r UniqueDictdirr startswithr#)prefixmapnamer0r(r(r)_build_reverse_error_maps   rZcCstS)N)r:getpidr(r(r(r)_getpidsr\z9driver missing function: %s. Requires CUDA 8.0 or above. c@seZdZdZdZddZddZddZd d Ze d d Z d dZ ddZ ddZ ddZd ddZddZddZddZddZdS)!Driverz0 Driver API functions are lazily bound. NcCs&|j}|dk r|St|}||_|S)N) _singletonobject__new__)clsobjr(r(r)r`s  zDriver.__new__c Csnt|_d|_d|_d|_ytjr0d}t|t |_ Wn.tk rh}zd|_||_Wdd}~XYnXdS)NFzzCUDA is disabled due to setting NUMBA_DISABLE_CUDA=1 in the environment, or because CUDA is unsupported on 32-bit systems.T) rrTdevicesis_initializedinitialization_errorpidrZ DISABLE_CUDArrQlib)r4r1rBr(r(r)r3s  zDriver.__init__c Csntad|_ytd|dWn4tk rX}z||_td|Wdd}~XYn Xt|_ | dS)NTZinitrzError at driver init: %s:) r*_loggerrdinfoZcuInitr/rerr\rf_initialize_extras)r4rBr(r(r) initializes zDriver.initializecCsltdt}|tj}||dttttj ttj tj }|tj }d|_ |d|}||_dS)NcuIpcOpenMemHandlecall_cuIpcOpenMemHandle)rHZ CFUNCTYPEr rset_cuIpcOpenMemHandle _find_apirZPOINTERr cu_device_ptrcu_ipc_mem_handleZc_uintrmr _wrap_api_callrl)r4Z set_protornZ call_protorm safe_callr(r(r)rjs     zDriver._initialize_extrascCs|js||jdkS)N)rdrkre)r4r(r(r) is_availableszDriver.is_availablecCsy t|}Wntk r(t|YnX|d}|dd}|jsL||jdk rdtd|j||}||_||_ | ||}t ||||S)NrrzError at driver init: %s:) rKeyErrorAttributeErrorrdrkrerrorestypeargtypesrrsetattr)r4fnameprotorwrxlibfnrsr(r(r) __getattr__s"       zDriver.__getattr__cstfdd}|S)Ncs&tdj|}|dS)Nzcall driver api: %s)rhdebugr _check_error)argsretcode)rzr|r4r(r)safe_cuda_api_callsz1Driver._wrap_api_call..safe_cuda_api_call) functoolswraps)r4rzr|rr()rzr|r4r)rrszDriver._wrap_api_callcsdyt|jdStk r$YnXy t|jStk rFYnXfdd}t|||S)NZ_v2csttdS)N)rMISSING_FUNCTION_ERRMSG)rZkws)rzr(r)absent_function4sz)Driver._find_api..absent_function)r#rgrvry)r4rzrr()rzr)ro%s   zDriver._find_apicCsx|tjkrtt|d}d||f}t||tjkrj|jdk rjt|jkrjd}t |t|jt dt ||dS)NZUNKNOWN_CUDA_ERRORzCall to %s results in %sz0pid %s forked from pid %s after CUDA driver initzCUDA initialized before forking) rZ CUDA_SUCCESS ERROR_MAPrErherrorZCUDA_ERROR_NOT_INITIALIZEDrfr\Zcriticalrr/)r4rzrZerrnamer1r(r(r)r:s     zDriver._check_errorrcCs0|j|}|dkr&t|}||j|<t|S)N)rcrEDeviceweakrefproxy)r4devnumdevr(r(r) get_deviceGs   zDriver.get_devicecCst}|t||jS)N)rZcuDeviceGetCountrvalue)r4countr(r(r)get_device_countNszDriver.get_device_countcCst|jS)z)Returns a list of active devices )listrcvalues)r4r(r(r) list_devicesSszDriver.list_devicescCs x|jD] }|q WdS)zReset all devices N)rcrreset)r4rr(r(r)rXsz Driver.resetcCs&td}tt||js"dS|S)zqGet current active context in CUDA driver runtime. Note: Lowlevel calls that returns the handle. rN)r cu_contextdriverZcuCtxGetCurrentrr)r4handler(r(r) get_context^s  zDriver.get_context)r)r r,r-__doc__r^r`r3rkrjpropertyrtr}rrrorrrrrrr(r(r(r)r]s    r]cCsFd}t}x4ttD](}||rtt|||t|d<qW|S)NZCU_DEVICE_ATTRIBUTE_)rrTrUrrVr#len)rWrXrYr(r(r)_build_reverse_device_attrsls  rc@sxeZdZdZeddZddZddZedd Z d d Z d d Z ddZ ddZ ddZddZddZddZdS)rz The device object owns the CUDA contexts. This is owned by the driver object. User should not construct devices directly. cCsFx@ttD]}t|}||kr|SqWd|}t|dS)zgCreate Device object from device identity created by ``Device.get_device_identity()``. zKNo device of {} is found. Target device may not be visible in this process.N)rangerrrget_device_identityformat RuntimeError)r4ZidentityZdeviddrPr(r(r) from_identity}s  zDevice.from_identitycCst}tt||||jks(td|j|_i|_t}t}tt|t||j|j|jf|_ d}t |}t |||j|j|_ d|_ dS)NzDriver returned another device)rrZ cuDeviceGetrrAssertionErroridZ attributesZcuDeviceComputeCapabilitycompute_capabilityr ZcuDeviceGetNamerYprimary_context)r4rZ got_devnumZcc_majorZcc_minorZbufszZbufr(r(r)r3s zDevice.__init__cCs|j|j|jdS)N)Z pci_domain_idZ pci_bus_idZ pci_device_id)Z PCI_DOMAIN_IDZ PCI_BUS_IDZ PCI_DEVICE_ID)r4r(r(r)rszDevice.get_device_identitycCstdt|jS)z, For backward compatibility zADeprecated attribute 'COMPUTE_CAPABILITY'; use lower case version)warningswarnDeprecationWarningr)r4r(r(r)COMPUTE_CAPABILITYszDevice.COMPUTE_CAPABILITYcCsd|j|jfS)Nz)rrY)r4r(r(r)__repr__szDevice.__repr__cCsXy t|}Wntk r(t|YnXt}tt|||jt|||j |j S)zRead attributes lazily ) DEVICE_ATTRIBUTESrurvrrZcuDeviceGetAttributerrryr)r4attrr0rr(r(r)r}s zDevice.__getattr__cCs t|jS)N)hashr)r4r(r(r)__hash__szDevice.__hash__cCst|tr|j|jkSdS)NF)r$rr)r4otherr(r(r)__eq__s  z Device.__eq__cCs ||k S)Nr()r4rr(r(r)__ne__sz Device.__ne__cCsL|jdk r|jSt|t}tt||jtt ||}||_|S)zo Returns the primary context for the device. Note: it is not pushed to the CPU thread. N) rmet_requirement_for_devicerrrZcuDevicePrimaryCtxRetainrrContextrr)r4ZhctxZctxr(r(r)get_primary_contexts zDevice.get_primary_contextcCst|jd|_dS)z6 Release reference to primary context N)rZcuDevicePrimaryCtxReleaserr)r4r(r(r)release_primary_contexts zDevice.release_primary_contextc Cs4z |jdk r|j|Wdt|jXdS)N)rrrrZcuDevicePrimaryCtxResetr)r4r(r(r)rs    z Device.resetN)r r,r-r classmethodrr3rrrrr}rrrrrrr(r(r(r)rxs  rcCs|jtkrtd|tfdS)Nz%s has compute capability < %s)rMIN_REQUIRED_CCr)devicer(r(r)rs rc@s eZdZdZddZddZdS) _SizeNotSetzC Dummy object for _PendingDeallocs when *size* is not set. cCsdS)N?r()r4r(r(r)r6sz_SizeNotSet.__str__cCsdS)Nrr()r4r(r(r)__int__sz_SizeNotSet.__int__N)r r,r-rr6rr(r(r(r)rsrc@sZeZdZdZddZeddZefddZdd Z e j d d Z ed d Z ddZdS)_PendingDeallocszd Pending deallocations of a context (or device since we are using the primary context). cCst|_d|_d|_||_dS)Nr)r_cons_disable_count_size_memory_capacity)r4Zcapacityr(r(r)r3sz_PendingDeallocs.__init__cCst|jtjS)N)r%rrZCUDA_DEALLOCS_RATIO)r4r(r(r)_max_pending_bytes sz#_PendingDeallocs._max_pending_bytescCs\td|j||j|||f|jt|7_t|jtj ksP|j|j krX| dS)a_ Add a pending deallocation. The *dtor* arg is the destructor function that takes an argument, *handle*. It is used as ``dtor(handle)``. The *size* arg is the byte size of the resource added. It is an optional argument. Some resources (e.g. CUModule) has an unknown memory footprint on the device. z add pending dealloc: %s %s bytesN) rhrir rrMrr%rrZCUDA_DEALLOCS_COUNTrclear)r4dtorrsizer(r(r)add_items  z_PendingDeallocs.add_itemcCsD|js@x2|jr8|j\}}}td|j|||qWd|_dS)zh Flush any pending deallocations unless it is disabled. Do nothing if disabled. zdealloc: %s %s bytesrN) is_disabledrpopleftrhrir r)r4rrrr(r(r)r!s  z_PendingDeallocs.clearc cs<|jd7_z dVWd|jd8_|jdks6tXdS)zs Context manager to temporarily disable flushing pending deallocation. This can be nested. rNr)rr)r4r(r(r)disable-s  z_PendingDeallocs.disablecCs |jdkS)Nr)r)r4r(r(r)r:sz_PendingDeallocs.is_disabledcCs t|jS)z: Returns number of pending deallocations. )rr)r4r(r(r)__len__>sz_PendingDeallocs.__len__N)r r,r-rr3rrrrr contextlibcontextmanagerrrrr(r(r(r)rs    r _MemoryInfoz free,totalc@seZdZdZddZddZddZd8d d Zd9d d Zd dZ ddZ ddZ ddZ d:ddZ d;ddZddZddZddZdrzs This object wraps a CUDA Context resource. Contexts should not be constructed directly by user code. cCs0||_||_t|_d|_t|_i|_dS)N)rrrrT allocations deallocationsmodulesZextras)r4rrr(r(r)r3Os   zContext.__init__cCsBtd|jjtd|jj|j|j|jdS)z? Clean up all owned resources in this context. zreset context of device %sN)rhrirrrrrr)r4r(r(r)rYs   z Context.resetcCs0t}t}tt|t|t|j|jdS)z>Returns (free, total) memory in bytes in the context. )freetotal)rrZ cuMemGetInforrr)r4rrr(r(r)get_memory_infoeszContext.get_memory_infoNcCs@t}|s"tt||j||ntt||j||||jS)zReturn occupancy of a function. :param func: kernel for which occupancy is calculated :param blocksize: block size the kernel is intended to be launched with :param memsize: per-block dynamic shared memory usage intended, in bytes)rrZ+cuOccupancyMaxActiveBlocksPerMultiprocessorrrZ4cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsr)r4func blocksizememsizeflagsretvalr(r(r)$get_active_blocks_per_multiprocessorms z,Context.get_active_blocks_per_multiprocessorc Csdt}t}t|}|s8tt|t||j|||n tt|t||j|||||j|jfS)aSuggest a launch configuration with reasonable occupancy. :param func: kernel for which occupancy is calculated :param b2d_func: function that calculates how much per-block dynamic shared memory 'func' uses based on the block size. :param memsize: per-block dynamic shared memory usage intended, in bytes :param blocksizelimit: maximum block size the kernel is designed to handle)rrrZ cuOccupancyMaxPotentialBlockSizerrZ)cuOccupancyMaxPotentialBlockSizeWithFlagsr) r4rZb2d_funcrZblocksizelimitrZgridsizerZb2d_cbr(r(r)get_max_potential_block_sizezs  z$Context.get_max_potential_block_sizecCs*t|j|jdkr&t|j|_dS)z@ Pushes this context on the current CPU Thread. N)rZcuCtxPushCurrentrrrrr)r4r(r(r)pushs  z Context.pushcCs,t}tt||j|jjks(tdS)z Pops this context off the current CPU thread. Note that this context must be at the top of the context stack, otherwise an error will occur. N)rrrZcuCtxPopCurrentrrrr)r4Zpoppedr(r(r)popsz Context.popc CsRy |WnBtk rL}z$|jtjkr:|j|nWdd}~XYnXdS)z Attempt allocation by calling *allocator*. If a out-of-memory error is raised, the pending deallocations are flushed and the allocation is retried. If it fails in the second attempt, the error is reraised. N)r/r0rZCUDA_ERROR_OUT_OF_MEMORYrr)r4 allocatorrBr(r(r)_attempt_allocations   zContext._attempt_allocationcsTtfdd}||t|}tt||}||jj<| S)NcsttdS)N)rZ cuMemAllocrr()bytesizeptrr(r)rsz#Context.memalloc..allocator) rrpr_alloc_finalizerAutoFreePointerrrrrown)r4rr finalizermemr()rrr)memallocs   zContext.memallocFc std|rtjO|r&tjO|r4tjOfdd}|rT||n|d}t||}|rtt |||d}||j |j j <| Stt |||d}|SdS)NrcsttdS)N)rZcuMemHostAllocrr()rrpointerr(r)rsz'Context.memhostalloc..allocator)r)r rZCU_MEMHOSTALLOC_DEVICEMAPZCU_MEMHOSTALLOC_PORTABLEZCU_MEMHOSTALLOC_WRITECOMBINEDr_hostalloc_finalizer MappedMemoryrrrrrr PinnedMemory) r4rmappedZportableZwcrownerrrr()rrrr) memhostallocs,    zContext.memhostalloccstttfrt|r0|jjs0td|jd|rBtjOfdd}|rb| |n|t ||}|rt t |||d}||j|jj<|Stt |||d}|SdS)Nz%s cannot map host memoryrcstdS)N)rZcuMemHostRegisterr()rrrr(r)rsz!Context.mempin..allocator)r)r$r%longr rZCAN_MAP_HOST_MEMORYrrZCU_MEMHOSTREGISTER_DEVICEMAPr_pin_finalizerrrrrrrrr)r4rrrrrrrr()rrrr)mempins(    zContext.mempincCstdS)N)NotImplementedError)r4rr(r(r)memunpinszContext.memunpincCsZts tdt}tt||jj |j }|j j |jj j }t |||j||dS)z> Returns a *IpcHandle* from a GPU allocation. zOS does not support CUDA IPC)offset) SUPPORTS_IPCrLrrqrZcuIpcGetMemHandlerHrrrrrr IpcHandler)r4ZmemoryZ ipchandle source_inforr(r(r)get_ipc_handles   zContext.get_ipc_handlecCs2t}d}tt|||tt|||dS)Nr)contextrr)rrprrlr MemoryPointerrr)r4rrdptrrr(r(r)open_ipc_handles  zContext.open_ipc_handlercCs |dkstdt||dS)zLEnable peer access between the current context and the peer context rz$*flags* is reserved and MUST be zeroN)rrZcuCtxEnablePeerAccess)r4Z peer_contextrr(r(r)enable_peer_accessszContext.enable_peer_accesscCs$t}tt||jj|t|S)zsReturns a bool indicating whether the peer access between the current and peer device is possible. )rrZcuDeviceCanAccessPeerrrrbool)r4Z peer_devicecan_access_peerr(r(r)r%s zContext.can_access_peercCs&t|tr|d}t|}||S)Nutf8)r$r!encoder create_module_image)r4ptximager(r(r)create_module_ptx1s  zContext.create_module_ptxcCs"t||}||j|jj<t|S)N)load_module_imagerrrrr)r4rmoduler(r(r)r7s zContext.create_module_imagecCs|j|jj=dS)N)rrr)r4rr(r(r) unload_module<szContext.unload_modulecCs2t}tt|dtt||t|j |S)Nr) rZ cu_streamrZcuStreamCreaterStreamrr_stream_finalizerr)r4rr(r(r) create_stream?s zContext.create_streamTcCsFt}d}|s|tjO}tt||tt ||t |j |dS)Nr)r) rZcu_eventrZCU_EVENT_DISABLE_TIMINGrZ cuEventCreaterEventrr_event_finalizerr)r4Ztimingrrr(r(r) create_eventEs  zContext.create_eventcCs tdS)N)rZcuCtxSynchronize)r4r(r(r) synchronizeNszContext.synchronizecCsd|j|jjfS)Nz)rrr)r4r(r(r)rQszContext.__repr__cCst|tr|j|jkStSdS)N)r$rrNotImplemented)r4rr(r(r)rTs  zContext.__eq__cCs || S)N)r)r4rr(r(r)rZszContext.__ne__)N)N)FFF)F)r)T)r r,r-rr3rrrrrrrrrrrrrrrrrr r rrrrrr(r(r(r)rHs4      #   rc Cs ttjdd}t|}t|}tjt|tjt |tj t|tj t |tj t t i}tjt||}t t||}t}ytt||t|||Wn>tk r} z d|jd} t| j| Wdd} ~ XYnX|j} tt||| t||S)z! image must be a pointer NUMBAPRO_CUDA_LOG_SIZEizcuModuleLoadDataEx error: %srN)r%r:rDrEr rCU_JIT_INFO_LOG_BUFFERr !CU_JIT_INFO_LOG_BUFFER_SIZE_BYTESr CU_JIT_ERROR_LOG_BUFFER"CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTESCU_JIT_LOG_VERBOSEVERBOSE_JIT_LOGr cu_jit_optionrkeysrZ cu_modulerZcuModuleLoadDataExrr/rdecoder0Modulerr_module_finalizer) rrlogszZjitinfoZ jiterrorsoptions option_keys option_valsrrBr1info_logr(r(r)r^s(       rcs"|j|jfdd}|S)Ncs r j=tjdS)N)rrrZ cuMemFreer()rrrrr(r)coresz_alloc_finalizer..core)rr)rrrr#r()rrrrr)rsrcs,|j|jstfdd}|S)a[ Finalize page-locked host memory allocated by `context.memhostalloc`. This memory is managed by CUDA, and finalization entails deallocation. The issues noted in `_pin_finalizer` are not relevant in this case, and the finalization is placed in the `context.deallocations` queue along with finalization of device objects. cs$rrj=tjdS)N)rrrZ cuMemFreeHostr()rrrrrr(r)r#sz"_hostalloc_finalizer..core)rrr)rrrrr#r()rrrrrr)rs rcs|jfdd}|S)aB Finalize temporary page-locking of host memory by `context.mempin`. This applies to memory not otherwise managed by CUDA. Page-locking can be requested multiple times on the same memory, and must therefore be lifted as soon as finalization is requested, otherwise subsequent calls to `mempin` may fail with `CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED`, leading to unexpected behavior for the context managers `cuda.{pinned,mapped}`. This function therefore carries out finalization immediately, bypassing the `context.deallocations` queue. csrrj=tdS)N)rrZcuMemHostUnregisterr()rrrr(r)r#sz_pin_finalizer..core)r)rrrr#r()rrrr)rs rcsfdd}|S)NcstjdS)N)rrZcuEventDestroyr()deallocsrr(r)r#sz_event_finalizer..corer()r$rr#r()r$rr)rsrcsfdd}|S)NcstjdS)N)rrZcuStreamDestroyr()r$rr(r)r#sz_stream_finalizer..corer()r$rr#r()r$rr)r sr cs |j|jfdd}|S)Ncs$tjfdd}|dS)Ncs"s|jkstt|dS)N)rrrZcuModuleUnload)r)r shutting_downr(r) module_unloadsz6_module_finalizer..core..module_unload)rr%r)r&)deallocrr)r%r)r#sz_module_finalizer..core)rr)rrr#r()r'rrr)rs rc@s(eZdZdZddZddZddZdS) _CudaIpcImplzjImplementation of GPU IPC using CUDA driver API. This requires the devices to be peer accessible. cCs*|j|_|j|_|j|_|j|_d|_dS)N)baserrr _opened_mem)r4parentr(r(r)r3s z_CudaIpcImpl.__init__cCsP|jdk rtd|jdk r$td||j|j|j}||_||jS)zT Import the IPC memory and returns a raw CUDA memory pointer object Nz'opening IpcHandle from original processzIpcHandle is already opened) r)rJr*rrrrrview)r4rrr(r(r)opens  z_CudaIpcImpl.opencCs*|jdkrtdt|jjd|_dS)NzIpcHandle not opened)r*rJrZcuIpcCloseMemHandler)r4r(r(r)closes z_CudaIpcImpl.closeN)r r,r-rr3r-r.r(r(r(r)r(sr(c@s(eZdZdZddZddZddZdS) _StagedIpcImplzImplementation of GPU IPC using custom staging logic to workaround CUDA IPC limitation on peer accessibility between devices. cCs(||_|j|_|j|_|j|_||_dS)N)r+r)rrr)r4r+rr(r(r)r3s z_StagedIpcImpl.__init__c Csddlm}t|j}t|jd}|j|j| |j }WdQRX| |j }t|||j |j|j|WdQRX|S)Nr)cuda)r+)numbar0rrrr(r+Zgpusrr-rcrrrdevice_to_devicer.r)r4rr0ZsrcdevimplZ source_ptrZnewmemr(r(r)r- s    z_StagedIpcImpl.opencCsdS)Nr()r4r(r(r)r. sz_StagedIpcImpl.closeN)r r,r-rr3r-r.r(r(r(r)r/sr/c@sheZdZdZdddZddZdd Zd d Zd d ZddZ dddZ ddZ ddZ e ddZdS)ra Internal IPC handle. Serialization of the CUDA IPC handle object is implemented here. The *base* attribute is a reference to the original allocation to keep it alive. The *handle* is a ctypes object of the CUDA IPC handle. The *size* is the allocation size. NrcCs(||_||_||_||_d|_||_dS)N)r)rrr_implr)r4r)rrrrr(r(r)r3/s zIpcHandle.__init__cCs|jdkrtddS)Nz#IPC handle doesn't have source info)rr)r4r(r(r)_sentry_source_info7s zIpcHandle._sentry_source_infocCs4||j|jkrdSt|j}||jS)zdReturns a bool indicating whether the active context can peer access the IPC handle T)r5rrrrrrr)r4rZ source_devicer(r(r)r;s  zIpcHandle.can_access_peercCs4||jdk rtdt||j|_|j|S)zCOpen the IPC by allowing staging on the host memory first. NzIpcHandle is already opened)r5r4rJr/rr-)r4rr(r(r) open_stagedEs  zIpcHandle.open_stagedcCs(|jdk rtdt||_|j|S)zT Import the IPC memory and returns a raw CUDA memory pointer object NzIpcHandle is already opened)r4rJr(r-)r4rr(r(r) open_directPs  zIpcHandle.open_directcCs*|jdks||r|j}n|j}||S)aOpen the IPC handle and import the memory for usage in the given context. Returns a raw CUDA memory pointer object. This is enhanced over CUDA IPC that it will work regardless of whether the source device is peer-accessible by the destination device. If the devices are peer-accessible, it uses .open_direct(). If the devices are not peer-accessible, it uses .open_staged(). N)rrr7r6)r4rfnr(r(r)r-Zs zIpcHandle.opencCs6ddlm}|dkr|j}||}|j||||dS)zC Simliar to `.open()` but returns an device array. r) devicearrayN)shapestridesdtypeZgpu_data)r9itemsizer-Z DeviceNDArray)r4rr:r<r;r9rr(r(r) open_arrayis   zIpcHandle.open_arraycCs&|jdkrtd|jd|_dS)NzIpcHandle not opened)r4rJr.)r4r(r(r)r.ws  zIpcHandle.closecCs*t|j}|j||j|j|jf}tj|fS)N)tuplerr5rrrrZ_rebuild_reduction)r4Zpreprocessed_handlerr(r(r) __reduce__}s zIpcHandle.__reduce__cCstj|}|d||||dS)N)r)rrrr)rrq)raZ handle_aryrrrrr(r(r)_rebuilds zIpcHandle._rebuild)Nr)N)r r,r-rr3r5rr6r7r-r?r.rArrBr(r(r(r)r%s       rc@sZeZdZdZdZdddZeddZdd Zd d Z dd dZ dddZ eddZ dS)raoA memory pointer that owns the buffer with an optional finalizer. When an instance is deleted, the finalizer will be called regardless of the `.refct`. An instance is created with `.refct=1`. The buffer lifetime is tied to the MemoryPointer instance's lifetime. The finalizer is invoked only if the MemoryPointer instance's lifetime ends. TNcCsP||_||_||_||_|dk |_d|_|j|_||_|dk rLt |||_ dS)Nr) rdevice_pointerr_cuda_memsize_ is_managedrefctr_ownerrfinalize _finalizer)r4rrrrrr(r(r)r3s zMemoryPointer.__init__cCs|jdkr|S|jS)N)rG)r4r(r(r)rszMemoryPointer.ownercCstt|S)N) OwnedPointerrr)r4r(r(r)rszMemoryPointer.owncCs.|jr*|jjstd||jjr*tdS)z8 Forces the device memory to the trash. zFreeing dead memoryN)rErIaliverr)r4r(r(r)rs zMemoryPointer.freercCs@|dkr|jn|}|r,t|j|||jnt|j||dS)N)rrcuMemsetD8AsyncrCr cuMemsetD8)r4Zbyterstreamr(r(r)memsets   zMemoryPointer.memsetcCs|dkr|j|}n||}|jjdkr>|dkr8td|}n:|jj|}|dkrZtdt|}t|j|||jd}t |jtt frt t |j|S|SdS)Nrz non-empty slice into empty slicezsize cannot be negative)r) rrCrrrrprrrr$rJrr)r4startstoprr,r)rr(r(r)r,s    zMemoryPointer.viewcCs|jS)N)rC)r4r(r(r)device_ctypes_pointersz#MemoryPointer.device_ctypes_pointer)NN)Nr)N) r r,r-r__cuda_memory__r3rrrrrOr,rRr(r(r(r)rs    rcs eZdZdZfddZZS)rzModifies the ownership semantic of the MemoryPointer so that the instance lifetime is directly tied to the number of references. When `.refct` reaches zero, the finalizer is invoked. cs$tt|j|||jd8_dS)Nr)r2rr3rF)r4rkwargs)r5r(r)r3szAutoFreePointer.__init__)r r,r-rr3r7r(r()r5r)rsrcs*eZdZdZdfdd ZddZZS)rTNcsb||_||_t}tt||d||_tt |j ||||d|j|_ |j |_ |jj|_dS)Nr)r)owned host_pointerrrprZcuMemHostGetDevicePointerrrCr2rr3rr_buflen_r_bufptr_)r4rrZ hostpointerrrdevptr)r5r(r)r3szMappedMemory.__init__cCstt|S)N)MappedOwnedPointerrr)r4r(r(r)rszMappedMemory.own)N)r r,r-rSr3rr7r(r()r5r)rsrc@seZdZdddZddZdS)rNcCsT||_||_||_||_|dk |_|j|_|j|_|jj|_|dk rPt ||dS)N) rrUrrVrErrWrrXrrH)r4rrrrrr(r(r)r3s  zPinnedMemory.__init__cCs|S)Nr()r4r(r(r)rszPinnedMemory.own)N)r r,r-r3rr(r(r(r)rs rc@seZdZdddZddZdS)rJNcsZ||_|dkr|j|_n|jr"t||_|jfdd}|jjd7_t||dS)NcsLy2jd8_jdkstjdkr0Wntk rFYnXdS)Nrr)rFrrReferenceErrorr()rr(r)deref$s  z$OwnedPointer.__init__..derefr)Z_mem_viewrErrFrrH)r4Zmemptrr,r\r()rr)r3s   zOwnedPointer.__init__cCs t|j|S)z$Proxy MemoryPointer methods )r#r])r4rzr(r(r)r}1szOwnedPointer.__getattr__)N)r r,r-r3r}r(r(r(r)rJs rJc@s eZdZdS)rZN)r r,r-r(r(r(r)rZ7srZc@s:eZdZddZddZddZddZejd d Z d S) r cCs$||_||_|dk r t||dS)N)rrrrH)r4rrrr(r(r)r3<szStream.__init__cCs|jjS)N)rr)r4r(r(r)rBszStream.__int__cCsd|jj|jfS)Nz)rrr)r4r(r(r)rEszStream.__repr__cCst|jdS)zy Wait for all commands in this stream to execute. This will commit any pending memory transfers. N)rZcuStreamSynchronizer)r4r(r(r)rHszStream.synchronizeccs|V|dS)z A context manager that waits for all commands in this stream to execute and commits any pending memory transfers upon exiting the context. N)r)r4r(r(r)auto_synchronizeOszStream.auto_synchronizeN) r r,r-r3rrrrrr^r(r(r(r)r ;s r c@sBeZdZdddZddZdddZd d Zdd d Zd dZdS)r NcCs$||_||_|dk r t||dS)N)rrrrH)r4rrrr(r(r)r3ZszEvent.__init__c CsNyt|jWn4tk rD}z|jtjkr2dSWdd}~XYnXdSdS)zy Returns True if all work before the most recent record has completed; otherwise, returns False. FNT)rZ cuEventQueryrr/r0rZCUDA_ERROR_NOT_READY)r4rBr(r(r)query`s z Event.queryrcCs |r |jnd}t|j|dS)a Set the record point of the event to the current point in the given stream. The event will be considered to have occurred when all work that was queued in the stream at the time of the call to ``record()`` has been completed. rN)rrZ cuEventRecord)r4rNhstreamr(r(r)recordos z Event.recordcCst|jdS)zN Synchronize the host thread for the completion of the event. N)rZcuEventSynchronizer)r4r(r(r)r{szEvent.synchronizecCs&|r |jnd}d}t||j|dS)zZ All future works submitted to stream will wait util the event completes. rN)rrZcuStreamWaitEvent)r4rNr`rr(r(r)waitsz Event.waitcCs t||S)N)event_elapsed_time)r4evtendr(r(r) elapsed_timeszEvent.elapsed_time)N)r)r) r r,r-r3r_rarrbrer(r(r(r)r Ys    r cCs"t}tt||j|j|jS)zF Compute the elapsed time between two events in milliseconds. )r rZcuEventElapsedTimerrr)ZevtstartrdZmsecr(r(r)rcsrcc@s.eZdZd ddZddZddZdd ZdS) rNcCs,||_||_||_|dk r(t|||_dS)N)rrr"rrHrI)r4rrr"rr(r(r)r3s zModule.__init__cCs|j|dS)N)rr )r4r(r(r)unloadsz Module.unloadcCs4t}tt||j|dtt |||S)Nr) rZ cu_functionrZcuModuleGetFunctionrrrFunctionrr)r4rYrr(r(r) get_functions zModule.get_functioncCsDt}t}tt|t||j|dt|j |||j fS)Nr) rrprrZcuModuleGetGlobalrrrrrr)r4rYrrr(r(r)get_global_symbols  zModule.get_global_symbol)N)r r,r-r3rfrhrir(r(r(r)rs rFuncAttrregssharedlocalconst maxthreadsc@sdeZdZdZdZdZdZddZddZddd Z dd d Z d d Z e ddZ ddZddZdS)rg)rrrrcCs ||_||_||_||_dS)N)rrrY_read_func_attr_allZattrs)r4rrrYr(r(r)r3szFunction.__init__cCs d|jS)Nz)rY)r4r(r(r)rszFunction.__repr__FcCsH|p |o |}|rtj}n|r$tj}n|r0tj}ntj}t|j|dS)N)rZCU_FUNC_CACHE_PREFER_EQUALZCU_FUNC_CACHE_PREFER_L1ZCU_FUNC_CACHE_PREFER_SHAREDZCU_FUNC_CACHE_PREFER_NONErZcuFuncSetCacheConfigr)r4Z prefer_equalZ prefer_cacheZ prefer_sharedflagr(r(r) cache_configs zFunction.cache_configcCsfxt|dkr|d7}qWxt|dkr2|d7}qWt|}||_||_||_|r\||_nd|_|S)N)rr)rcopygriddimblockdim sharedmemrN)r4rurvrwrNZinstr(r(r) configures   zFunction.configurecGs2|jr|jj}nd}t|j|j|j|j||dS)zS *args -- Must be either ctype objects of DevicePointer instances. N)rNr launch_kernelrurvrw)r4rZ streamhandler(r(r)__call__s  zFunction.__call__cCs |jjjS)N)rrr)r4r(r(r)rszFunction.devicecCs t}tt|||j|jS)z, Read CUfunction attributes )rrZcuFuncGetAttributerrr)r4Zattridrr(r(r)_read_func_attrszFunction._read_func_attrcCsN|tj}|tj}|tj}|tj}|tj}t|||||dS)N)rkrnrmrlro)r{rZCU_FUNC_ATTRIBUTE_NUM_REGSZ"CU_FUNC_ATTRIBUTE_CONST_SIZE_BYTESZ"CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTESZ#CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTESZ'CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCKrj)r4ZnregsZcmemZlmemZsmemZmaxtpbr(r(r)rps     zFunction._read_func_attr_allN)FFF)rr)r r,r-rurvrNrwr3rrrrxrzrrr{rpr(r(r(r)rgs  rgc Cs|\}}}|\} } } g} x6|D].} t| r>| tt| q| t| qWtt| | }t||||| | | |||d dS)N)is_device_memoryrMr rRr rrZcuLaunchKernel)Z cufunc_handlerurvrwr`rZgxZgyZgzZbxZbyZbzZ param_valsargZparamsr(r(r)rys   ry)oracubinZfatbinc@sPeZdZdddZeddZeddZdd d Zd d Zd dZ ddZ dS)Linkerrc Csttjdd}t|}t|}tjt|tjt |tj t|tj t |tj t di}|rnt ||tj <t|tjg}t|}~tjt||}t t||} t|_} tt||| t|jt|tj| ||_||_|||| g|_dS)Nrir) r%r:rDrEr rrr rr rrrZCU_JIT_MAX_REGISTERSrrZCU_JIT_TARGET_FROM_CUCONTEXTrrrrZ cu_link_staterrZ cuLinkCreaterrrHZ cuLinkDestroylinker_info_buflinker_errors_buf _keep_alive) r4Z max_registersrZ linkerinfoZ linkererrorsrZraw_keysZ raw_valuesr r!rr(r(r)r3#s,        zLinker.__init__cCs|jjdS)Nr)rrr)r4r(r(r)r"DszLinker.info_logcCs|jjdS)Nr)rrr)r4r(r(r) error_logHszLinker.error_log c Cst|}t|d}|j||g7_y$t|jtj|t||dddWn4t k r}zt d||j fWdd}~XYnXdS)Nrrz%s %s) r rrrZ cuLinkAddDatarrCU_JIT_INPUT_PTXrr/r.r)r4rrYZptxbufZnamebufrBr(r(r)add_ptxLs zLinker.add_ptxc Csnt|d}|j|yt|j||dddWn4tk rh}ztd||j fWdd}~XYnXdS)Nrrz%s %s) r rrrMrZ cuLinkAddFilerr/r.r)r4r;kindZpathbufrBr(r(r)add_fileVs  zLinker.add_filecCs(|ddd}t|}|||dS)N.r)rsplitFILE_EXTENSION_MAPr)r4r;Zextrr(r(r)add_file_guess_ext_szLinker.add_file_guess_extc Cstd}td}yt|jt|t|Wn4tk r`}ztd||jfWdd}~XYnX|j }|dksxt d|j dd=||fS)z Returns (cubin, size) cubin is a pointer to a internal buffer of cubin owned by the linker; thus, it should be loaded before the linker is destroyed. rz%s %sNz"linker returned a zero sized cubin) r rrZcuLinkCompleterrr/r.rrrr)r4rrrBr(r(r)completeds$ zLinker.completeN)r)r) r r,r-r3rr"rrrrrr(r(r(r)r"s !    rcCs&tt||t|}t|ddS)z*Query attribute on the device pointer z!Failed to query pointer attributeN)rcuPointerGetAttributerrRZ check_error)devmemrZodatarr(r(r)_device_pointer_attr|s  rcCs<td}t|tj|tjdtjdtjdtjdi}||jS)zAQuery the device pointer type: host, device, array, unified? rZhostrZarrayZunified) rrrZ CU_POINTER_ATTRIBUTE_MEMORY_TYPEZCU_MEMORYTYPE_HOSTZCU_MEMORYTYPE_DEVICEZCU_MEMORYTYPE_ARRAYZCU_MEMORYTYPE_UNIFIEDr)rZptrtyperXr(r(r)device_pointer_types rcCs$td}tj}tt||||S)zZQuery the device pointer usable in the current context from an arbitrary pointer. r)r rZ#CU_POINTER_ATTRIBUTE_DEVICE_POINTERrrr)rrYrr(r(r)get_devptr_for_active_ctxsrcCsFt}t}t|}tt|t|||j|j}}|||fS)aFind the extents (half open begin and end pointer) of the underlying device memory allocation. NOTE: it always returns the extents of the allocation but the extents of the device memory view that can be a subsection of the entire allocation. )rrprrRrZcuMemGetAddressRangerr)rsnrYr(r(r)device_extentss rcCsHt|dd}|dkr.t|\}}||}||_|dksDtd||S)zCheck the memory size of the device memory. The result is cached in the device memory object. It may query the driver for the memory size of the device memory allocation. rDNrz{} length array)r#rrDrr)rZszrrBr(r(r)device_memory_sizes  rcCst|dd}|dk o|jdkS)z?Returns True if the obj.dtype is datetime64 or timedelta64 r<NZMm)r#char)rbr<r(r(r)_is_datetime_dtypes rcCst|r|tj}|S)z^Workaround for numpy#4983: buffer protocol doesn't support datetime64 or timedelta64. )rr,npZint64)rbr(r(r)_workaround_for_datetimes rFcCsDt|ttfr|Sd}|s.t|tjp,t|}t|}t|||S)aGet host pointer from an obj. If `readonly` is False, the buffer must be writable. NOTE: The underlying data pointer from the host data buffer is used and it should not be changed until the operation which can be asynchronous completes. F) r$r%rrZvoidrrrZmemoryview_get_buffer)rbreadonlyZ forcewritabler(r(r)rVs rVcCst|}t|S)zHReturns (start, end) the start and end pointer of the array (half open).)rrZmemoryview_get_extents)rbr(r(r)host_memory_extentssrcCs<t|t|kstdt|}t||||\}}||S)z_Get the byte size of a contiguous memory buffer given the shape, strides and itemsize. z# dim mismatch)rrrZmemoryview_get_extents_info)r:r;r>ndimrrBr(r(r)memory_size_from_infosrcCs$t|\}}||kstd||S)zGet the size of the memoryzmemory extend of negative size)rr)rbrrBr(r(r)host_memory_sizes rcCs t|jS)z$Get the device pointer as an integer)rRr)rbr(r(r)rCsrCcCs|dkrtdSt||jS)z,Get the ctypes object for the device pointerNr)r require_device_memoryrR)rbr(r(r)rRsrRcCs t|ddS)ahAll CUDA memory object is recognized as an instance with the attribute "__cuda_memory__" defined and its value evaluated to True. All CUDA memory object should also define an attribute named "device_pointer" which value is an int(or long) object carrying the pointer value of the device memory address. This is not tested in this method. rSF)r#)rbr(r(r)r|sr|cCst|stddS)z9A sentry for methods that accept CUDA memory object. zNot a CUDA memory object.N)r| Exception)rbr(r(r)r srcGst|dg}||dS)zAdd dependencies to the device memory. Mainly used for creating structures that points to other device memory, so that the referees are not GC and released. Z _depends_N)r#extend)rZobjsZdepsetr(r(r)device_memory_dependss rcCsRg}|r*t|tsttj}||jntj}|t|t |dd|f|dS)z NOTE: The underlying data pointer from the host data buffer is used and it should not be changed until the operation which can be asynchronous completes. T)rN) r$r rrZcuMemcpyHtoDAsyncrMrZ cuMemcpyHtoDrCrV)dstsrcrrNvarargsr8r(r(r)host_to_devicesrcCsNg}|r*t|tsttj}||jntj}|t|t ||f|dS)z NOTE: The underlying data pointer from the host data buffer is used and it should not be changed until the operation which can be asynchronous completes. N) r$r rrZcuMemcpyDtoHAsyncrMrZ cuMemcpyDtoHrVrC)rrrrNrr8r(r(r)device_to_host.srcCsNg}|r*t|tsttj}||jntj}|t|t||f|dS)z NOTE: The underlying data pointer from the host data buffer is used and it should not be changed until the operation which can be asynchronous completes. N) r$r rrZcuMemcpyDtoDAsyncrMrZ cuMemcpyDtoDrC)rrrrNrr8r(r(r)r2@sr2cCsJg}|r*t|tsttj}||jntj}|t|||f|dS)zMemset on the device. If stream is not zero, asynchronous mode is used. dst: device memory val: byte value to be written size: number of byte to be written stream: a CUDA stream N) r$r rrrLrMrrMrC)rvalrrNrr8r(r(r) device_memsetRs rcCs tdS)z; Enable profile collection in the current context. N)rZcuProfilerStartr(r(r(r) profile_startgsrcCs tdS)z< Disable profile collection in the current context. N)rZcuProfilerStopr(r(r(r) profile_stopnsrccstdVtdS)z] Context manager that enables profiling on entry and disables profiling on exit. N)rrr(r(r(r) profilingusr)F)r)r)r)r)rZ __future__rrrr&r:rHrrrtrr itertoolsrrrrr r r r r rZnumpyr collectionsrrr1rrrrrrrrr=rrrrZ numba.utilsrrr%rDrErrrGrVrr*rr+r.r/rQrRrSrFrOrZr\rrr_r]rrrrrrrrrrrrrrr rr(r/rrrrZMemAllocrrJrZr r rcrrjrgryZCU_JIT_INPUT_OBJECTrZCU_JIT_INPUT_LIBRARYZCU_JIT_INPUT_CUBINZCU_JIT_INPUT_FATBINARrrrrrrrrrrVrrrrCrRr|rrrrr2rrrrrr(r(r(r) s (     A . w C " $&kQ 4  P Z