drvapi.py

from ctypes import (c_byte, c_char_p, c_float, c_int, c_size_t, c_uint,
                    c_uint8, c_void_p, py_object, CFUNCTYPE, POINTER)

from numba.cuda.cudadrv import _extras

cu_device = c_int
cu_device_attribute = c_int     # enum
cu_context = c_void_p           # an opaque handle
cu_module = c_void_p            # an opaque handle
cu_jit_option = c_int           # enum
cu_jit_input_type = c_int       # enum
cu_function = c_void_p          # an opaque handle
cu_device_ptr = c_size_t        # defined as unsigned long long
cu_stream = c_void_p            # an opaque handle
cu_event = c_void_p
cu_link_state = c_void_p
cu_function_attribute = c_int
cu_ipc_mem_handle = (c_byte * _extras.CUDA_IPC_HANDLE_SIZE)   # 64 bytes wide
cu_uuid = (c_byte * 16)         # Device UUID

cu_stream_callback_pyobj = CFUNCTYPE(None, cu_stream, c_int, py_object)

cu_occupancy_b2d_size = CFUNCTYPE(c_size_t, c_int)

# See https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TYPES.html
CU_STREAM_DEFAULT = 0
CU_STREAM_LEGACY = 1
CU_STREAM_PER_THREAD = 2

API_PROTOTYPES = {
    # CUresult cuInit(unsigned int Flags);
    'cuInit' : (c_int, c_uint),

    # CUresult cuDriverGetVersion (int* driverVersion )
    'cuDriverGetVersion': (c_int, POINTER(c_int)),

    # CUresult cuDeviceGetCount(int *count);
    'cuDeviceGetCount': (c_int, POINTER(c_int)),

    # CUresult cuDeviceGet(CUdevice *device, int ordinal);
    'cuDeviceGet': (c_int, POINTER(cu_device), c_int),

    # CUresult cuDeviceGetName ( char* name, int  len, CUdevice dev )
    'cuDeviceGetName': (c_int, c_char_p, c_int, cu_device),

    # CUresult cuDeviceGetAttribute(int *pi, CUdevice_attribute attrib,
    #                               CUdevice dev);
    'cuDeviceGetAttribute': (c_int, POINTER(c_int), cu_device_attribute,
                             cu_device),

    # CUresult cuDeviceComputeCapability(int *major, int *minor,
    #                                    CUdevice dev);
    'cuDeviceComputeCapability': (c_int, POINTER(c_int), POINTER(c_int),
                                  cu_device),

    # CUresult cuDevicePrimaryCtxGetState(
    #              CUdevice dev,
    #              unsigned int* flags,
    #              int* active)
    'cuDevicePrimaryCtxGetState': (c_int,
                                   cu_device, POINTER(c_uint), POINTER(c_int)),

    # CUresult cuDevicePrimaryCtxRelease ( CUdevice dev )
    'cuDevicePrimaryCtxRelease': (c_int, cu_device),

    # CUresult cuDevicePrimaryCtxReset ( CUdevice dev )
    'cuDevicePrimaryCtxReset': (c_int, cu_device),

    # CUresult cuDevicePrimaryCtxRetain ( CUcontext* pctx, CUdevice dev )
    'cuDevicePrimaryCtxRetain': (c_int, POINTER(cu_context), cu_device),

    # CUresult cuDevicePrimaryCtxSetFlags ( CUdevice dev, unsigned int  flags )
    'cuDevicePrimaryCtxSetFlags': (c_int, cu_device, c_uint),

    # CUresult cuCtxCreate(CUcontext *pctx, unsigned int flags,
    #                      CUdevice dev);
    'cuCtxCreate': (c_int, POINTER(cu_context), c_uint, cu_device),

    # CUresult cuCtxGetDevice (	CUdevice * 	device	 )
    'cuCtxGetDevice': (c_int, POINTER(cu_device)),

    # CUresult cuCtxGetCurrent (CUcontext *pctx);
    'cuCtxGetCurrent': (c_int, POINTER(cu_context)),

    # CUresult cuCtxPushCurrent (CUcontext pctx);
    'cuCtxPushCurrent': (c_int, cu_context),

    # CUresult cuCtxPopCurrent (CUcontext *pctx);
    'cuCtxPopCurrent': (c_int, POINTER(cu_context)),

    # CUresult cuCtxDestroy(CUcontext pctx);
    'cuCtxDestroy': (c_int, cu_context),

    # CUresult cuModuleLoadDataEx(CUmodule *module, const void *image,
    #                             unsigned int numOptions,
    #                             CUjit_option *options,
    #                             void **optionValues);
    'cuModuleLoadDataEx': (c_int, cu_module, c_void_p, c_uint,
                           POINTER(cu_jit_option), POINTER(c_void_p)),

    # CUresult cuModuleUnload(CUmodule hmod);
    'cuModuleUnload': (c_int, cu_module),

    # CUresult cuModuleGetFunction(CUfunction *hfunc, CUmodule hmod,
    #                              const char *name);
    'cuModuleGetFunction': (c_int, cu_function, cu_module, c_char_p),

    # CUresult cuModuleGetGlobal ( CUdeviceptr* dptr, size_t* bytes, CUmodule
    #                              hmod, const char* name )
    'cuModuleGetGlobal': (c_int, POINTER(cu_device_ptr), POINTER(c_size_t),
                          cu_module, c_char_p),

    # CUresult CUDAAPI cuFuncSetCacheConfig(CUfunction hfunc,
    #                                       CUfunc_cache config);
    'cuFuncSetCacheConfig': (c_int, cu_function, c_uint),

    # CUresult cuMemAlloc(CUdeviceptr *dptr, size_t bytesize);
    'cuMemAlloc': (c_int, POINTER(cu_device_ptr), c_size_t),

    # CUresult cuMemAllocManaged(CUdeviceptr *dptr, size_t bytesize,
    #                            unsigned int flags);
    'cuMemAllocManaged': (c_int, c_void_p, c_size_t, c_uint),

    # CUresult cuMemsetD8(CUdeviceptr dstDevice, unsigned char uc, size_t N)
    'cuMemsetD8': (c_int, cu_device_ptr, c_uint8, c_size_t),

    # CUresult cuMemsetD8Async(CUdeviceptr dstDevice, unsigned char uc,
    #                          size_t N, CUstream hStream);
    'cuMemsetD8Async': (c_int,
                        cu_device_ptr, c_uint8, c_size_t, cu_stream),

    # CUresult cuMemcpyHtoD(CUdeviceptr dstDevice, const void *srcHost,
    #                       size_t ByteCount);
    'cuMemcpyHtoD': (c_int, cu_device_ptr, c_void_p, c_size_t),

    # CUresult cuMemcpyHtoDAsync(CUdeviceptr dstDevice, const void *srcHost,
    #                            size_t ByteCount, CUstream hStream);
    'cuMemcpyHtoDAsync': (c_int, cu_device_ptr, c_void_p, c_size_t,
                          cu_stream),

    # CUresult cuMemcpyDtoD(CUdeviceptr dstDevice, const void *srcDevice,
    #                       size_t ByteCount);
    'cuMemcpyDtoD': (c_int, cu_device_ptr, cu_device_ptr, c_size_t),

    # CUresult cuMemcpyDtoDAsync(CUdeviceptr dstDevice, const void *srcDevice,
    #                            size_t ByteCount, CUstream hStream);
    'cuMemcpyDtoDAsync': (c_int, cu_device_ptr, cu_device_ptr, c_size_t,
                          cu_stream),


    # CUresult cuMemcpyDtoH(void *dstHost, CUdeviceptr srcDevice,
    #                       size_t ByteCount);
    'cuMemcpyDtoH': (c_int, c_void_p, cu_device_ptr, c_size_t),

    # CUresult cuMemcpyDtoHAsync(void *dstHost, CUdeviceptr srcDevice,
    #                            size_t ByteCount, CUstream hStream);
    'cuMemcpyDtoHAsync': (c_int, c_void_p, cu_device_ptr, c_size_t,
                          cu_stream),

    # CUresult cuMemFree(CUdeviceptr dptr);
    'cuMemFree': (c_int, cu_device_ptr),

    # CUresult cuStreamCreate(CUstream *phStream, unsigned int Flags);
    'cuStreamCreate': (c_int, POINTER(cu_stream), c_uint),

    # CUresult cuStreamDestroy(CUstream hStream);
    'cuStreamDestroy': (c_int, cu_stream),

    # CUresult cuStreamSynchronize(CUstream hStream);
    'cuStreamSynchronize': (c_int, cu_stream),

    # CUresult cuStreamAddCallback(
    #              CUstream hStream,
    #              CUstreamCallback callback,
    #              void* userData,
    #              unsigned int flags)
    'cuStreamAddCallback': (c_int, cu_stream, cu_stream_callback_pyobj,
                            py_object, c_uint),

    # CUresult cuLaunchKernel(CUfunction f, unsigned int gridDimX,
    #                        unsigned int gridDimY,
    #                        unsigned int gridDimZ,
    #                        unsigned int blockDimX,
    #                        unsigned int blockDimY,
    #                        unsigned int blockDimZ,
    #                        unsigned int sharedMemBytes,
    #                        CUstream hStream, void **kernelParams,
    #                        void ** extra)
    'cuLaunchKernel': (c_int, cu_function, c_uint, c_uint, c_uint,
                       c_uint, c_uint, c_uint, c_uint, cu_stream,
                       POINTER(c_void_p), POINTER(c_void_p)),

    # CUresult cuLaunchCooperativeKernel(CUfunction f, unsigned int gridDimX,
    #                                   unsigned int gridDimY,
    #                                   unsigned int gridDimZ,
    #                                   unsigned int blockDimX,
    #                                   unsigned int blockDimY,
    #                                   unsigned int blockDimZ,
    #                                   unsigned int sharedMemBytes,
    #                                   CUstream hStream, void **kernelParams)
    'cuLaunchCooperativeKernel': (c_int, cu_function, c_uint, c_uint, c_uint,
                                  c_uint, c_uint, c_uint, c_uint, cu_stream,
                                  POINTER(c_void_p)),

    #  CUresult cuMemHostAlloc (	void ** 	pp,
    #                               size_t 	bytesize,
    #                               unsigned int 	Flags
    #                           )
    'cuMemHostAlloc': (c_int, c_void_p, c_size_t, c_uint),

    #  CUresult cuMemFreeHost (	void * 	p	 )
    'cuMemFreeHost': (c_int, c_void_p),

    # CUresult cuMemHostRegister(void * 	p,
    #                            size_t 	bytesize,
    #                            unsigned int 	Flags)
    'cuMemHostRegister': (c_int, c_void_p, c_size_t, c_uint),

    # CUresult cuMemHostUnregister(void * 	p)
    'cuMemHostUnregister': (c_int, c_void_p),

    # CUresult cuMemHostGetDevicePointer(CUdeviceptr * pdptr,
    #                                    void *        p,
    #                                    unsigned int  Flags)
    'cuMemHostGetDevicePointer': (c_int, POINTER(cu_device_ptr),
                                  c_void_p, c_uint),

    # CUresult cuMemGetInfo(size_t * free, size_t * total)
    'cuMemGetInfo' : (c_int, POINTER(c_size_t), POINTER(c_size_t)),

    # CUresult cuEventCreate (	CUevent * 	phEvent,
    #                               unsigned int 	Flags )
    'cuEventCreate': (c_int, POINTER(cu_event), c_uint),

    # CUresult cuEventDestroy (	CUevent 	hEvent	 )
    'cuEventDestroy': (c_int, cu_event),

    # CUresult cuEventElapsedTime (	float * 	pMilliseconds,
    #                                   CUevent 	hStart,
    #                                   CUevent 	hEnd )
    'cuEventElapsedTime': (c_int, POINTER(c_float), cu_event, cu_event),

    # CUresult cuEventQuery (	CUevent 	hEvent	 )
    'cuEventQuery': (c_int, cu_event),

    # CUresult cuEventRecord (	CUevent 	hEvent,
    #                               CUstream 	hStream )
    'cuEventRecord': (c_int, cu_event, cu_stream),

    # CUresult cuEventSynchronize (	CUevent 	hEvent	 )
    'cuEventSynchronize': (c_int, cu_event),


    # CUresult cuStreamWaitEvent (	CUstream        hStream,
    #                                   CUevent         hEvent,
    #                                	unsigned int 	Flags )
    'cuStreamWaitEvent': (c_int, cu_stream, cu_event, c_uint),

    # CUresult 	cuPointerGetAttribute (
    #               void *data,
    #               CUpointer_attribute attribute,
    #               CUdeviceptr ptr)
    'cuPointerGetAttribute': (c_int, c_void_p, c_uint, cu_device_ptr),

    #    CUresult cuMemGetAddressRange (	CUdeviceptr * 	pbase,
    #                                        size_t * 	psize,
    #                                        CUdeviceptr 	dptr
    #                                        )
    'cuMemGetAddressRange': (c_int,
                             POINTER(cu_device_ptr),
                             POINTER(c_size_t),
                             cu_device_ptr),

    #    CUresult cuMemHostGetFlags (	unsigned int * 	pFlags,
    #                                   void * 	p )
    'cuMemHostGetFlags': (c_int,
                          POINTER(c_uint),
                          c_void_p),

    #   CUresult cuCtxSynchronize ( void )
    'cuCtxSynchronize' : (c_int,),

    #    CUresult
    #    cuLinkCreate(unsigned int numOptions, CUjit_option *options,
    #                 void **optionValues, CUlinkState *stateOut);
    'cuLinkCreate': (c_int,
                     c_uint, POINTER(cu_jit_option),
                     POINTER(c_void_p), POINTER(cu_link_state)),

    #    CUresult
    #    cuLinkAddData(CUlinkState state, CUjitInputType type, void *data,
    #                  size_t size, const char *name, unsigned
    #                  int numOptions, CUjit_option *options,
    #                  void **optionValues);
    'cuLinkAddData': (c_int,
                      cu_link_state, cu_jit_input_type, c_void_p,
                      c_size_t, c_char_p, c_uint, POINTER(cu_jit_option),
                      POINTER(c_void_p)),

    #    CUresult
    #    cuLinkAddFile(CUlinkState state, CUjitInputType type,
    #                  const char *path, unsigned int numOptions,
    #                  CUjit_option *options, void **optionValues);

    'cuLinkAddFile': (c_int,
                      cu_link_state, cu_jit_input_type, c_char_p, c_uint,
                      POINTER(cu_jit_option), POINTER(c_void_p)),

    #    CUresult CUDAAPI
    #    cuLinkComplete(CUlinkState state, void **cubinOut, size_t *sizeOut)
    'cuLinkComplete': (c_int,
                       cu_link_state, POINTER(c_void_p), POINTER(c_size_t)),

    #    CUresult CUDAAPI
    #    cuLinkDestroy(CUlinkState state)
    'cuLinkDestroy': (c_int, cu_link_state),

    # cuProfilerStart ( void )
    'cuProfilerStart': (c_int,),

    # cuProfilerStop ( void )
    'cuProfilerStop': (c_int,),

    # CUresult cuFuncGetAttribute ( int* pi, CUfunction_attribute attrib,
    #                              CUfunction hfunc )
    'cuFuncGetAttribute': (c_int,
                           POINTER(c_int), cu_function_attribute, cu_function),

    # CUresult CUDAAPI cuOccupancyMaxActiveBlocksPerMultiprocessor(
    #                      int *numBlocks,
    #                      CUfunction func,
    #                      int blockSize,
    #                      size_t dynamicSMemSize);
    'cuOccupancyMaxActiveBlocksPerMultiprocessor': (c_int, POINTER(c_int),
                                                    cu_function, c_size_t,
                                                    c_uint),

    # CUresult CUDAAPI cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
    #                      int *numBlocks,
    #                      CUfunction func,
    #                      int blockSize,
    #                      size_t dynamicSMemSize,
    #                      unsigned int flags);
    'cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlags': (c_int,
                                                             POINTER(c_int),
                                                             cu_function,
                                                             c_size_t, c_uint),

    # CUresult CUDAAPI cuOccupancyMaxPotentialBlockSize(
    #                      int *minGridSize, int *blockSize,
    #                      CUfunction func,
    #                      CUoccupancyB2DSize blockSizeToDynamicSMemSize,
    #                      size_t dynamicSMemSize, int blockSizeLimit);
    'cuOccupancyMaxPotentialBlockSize': (c_int, POINTER(c_int), POINTER(c_int),
                                         cu_function, cu_occupancy_b2d_size,
                                         c_size_t, c_int),

    # CUresult CUDAAPI cuOccupancyMaxPotentialBlockSizeWithFlags(
    #                      int *minGridSize, int *blockSize,
    #                      CUfunction func,
    #                      CUoccupancyB2DSize blockSizeToDynamicSMemSize,
    #                      size_t dynamicSMemSize, int blockSizeLimit,
    #                      unsigned int flags);
    'cuOccupancyMaxPotentialBlockSizeWithFlags': (c_int, POINTER(c_int),
                                                  POINTER(c_int), cu_function,
                                                  cu_occupancy_b2d_size,
                                                  c_size_t, c_int, c_uint),

    # CUresult cuIpcGetMemHandle ( CUipcMemHandle* pHandle, CUdeviceptr dptr )
    'cuIpcGetMemHandle': (c_int,
                          POINTER(cu_ipc_mem_handle), cu_device_ptr),

    # CUresult cuIpcOpenMemHandle(
    #              CUdeviceptr* pdptr,
    #              CUipcMemHandle handle,
    #              unsigned int Flags)
    'cuIpcOpenMemHandle': (c_int, POINTER(cu_device_ptr), cu_ipc_mem_handle,
                           c_uint),

    # CUresult cuIpcCloseMemHandle ( CUdeviceptr dptr )

    'cuIpcCloseMemHandle': (c_int, cu_device_ptr),

    # CUresult cuCtxEnablePeerAccess (CUcontext peerContext, unsigned int Flags)
    'cuCtxEnablePeerAccess': (c_int, cu_context, c_int),

    # CUresult cuDeviceCanAccessPeer ( int* canAccessPeer,
    #                                  CUdevice dev, CUdevice peerDev )
    'cuDeviceCanAccessPeer': (c_int,
                              POINTER(c_int), cu_device, cu_device),

    # CUresult cuDeviceGetUuid ( CUuuid* uuid, CUdevice dev )
    'cuDeviceGetUuid': (c_int, POINTER(cu_uuid), cu_device),
}