o n hM@s,ddlZddlZddlZddlZddlZddlZddlmZddlm Z ddl m Z ddl m Z ddlmZddlmZejejeZejedgZejed Zd gZed d Zed dZeddZddZGdddeZ ddZ!ddZ"GdddeZ#GdddeZ$dS)N)Path)_build)get_cache_manager) _allocation) GPUTarget) GPUDriverincludelibcudacCstd}|r |gStddg}dd|D}dd|D}td}|r6|s6dd|d D}d }|rG|d t|7}|d 7}n|d 7}|d7}tdd|Ds\J||S)NTRITON_LIBCUDA_PATHz/sbin/ldconfigz-pcSs g|] }d|vr|dqS) libcuda.so.1)split).0linerq/var/www/html/construction_image-detection-poc/venv/lib/python3.10/site-packages/triton/backends/nvidia/driver.py s z libcuda_dirs..cSsg|]}tj|qSr)ospathdirname)rlocrrrrsLD_LIBRARY_PATHcSs&g|]}tjtj|dr|qS)r rrexistsjoin)rdirrrrr!s&:zlibcuda.so cannot found! z!Possible files are located at %s.z:Please create a symlink of libcuda.so to any of the files.z)s$zlibcuda_dirs..) rgetenv subprocess check_outputdecode splitlinesrstrany)env_libcuda_pathlibslocsdirsenv_ld_library_pathmsgrrr libcuda_dirss"   r,cCs tgtSN) libdevice_dirr,rrrr library_dirs-s r/cCs,ddlm}m}m}d||g|S)Nrmachinesystem architecture,)platformr1r2r3rr0rrr platform_key2sr6c CsJt|td}t|}tddd}| |d|}|durt W}t j |d}t|d }||Wdn1sJwYt|||ttt} t| d}|j||d|dd }Wdn1sxwYWdn1swYd dl} | j||} | j| } | j| | S) Nzutf-8 EXT_SUFFIX.r zmain.cwrbT)binaryr)hashlibsha256r6encode hexdigestr sysconfigget_config_varrget_filetempfileTemporaryDirectoryrrropenwriterr/ include_dir librariesputreadimportlib.utilutilspec_from_file_locationmodule_from_specloader exec_module) srcnamekeycacheext cache_pathtmpdirsrc_pathfso importlibspecmodrrrcompile_module_from_src8s*       r^cs$eZdZfddZddZZS) CudaUtilscs"t|dstt|||_|jS)Ninstance)hasattrsuperr___new__r`)cls __class__rrrcSs zCudaUtils.__new__cCsPtttjtdd}|j|_|j|_|j |_ |j |_ |j |_ |j |_ dS)Nzdriver.c cuda_utils) r^rrrrr read_text load_binaryget_device_propertiescuOccupancyMaxActiveClustersset_printf_fifo_sizefill_1d_tma_descriptorfill_2d_tma_descriptor)selfr]rrr__init__Xs zCudaUtils.__init__)__name__ __module__ __qualname__rcrp __classcell__rrrerr_Qs r_cCsx|ddkrdSidddddd d dd d d ddddddddddddddddddddd|S)Nr* CUdeviceptri1int32_ti8int8_ti16int16_ti32i64int64_tu1uint32_tu8uint8_tu16uint16_tu32u64uint64_tfp16floatbf16fp32f32fp64double nvTmaDesc CUtensorMapr)tyrrr ty_to_cppgsJ      rcsfddfddfdddfdd |D}d |}d t|}ttt|d }d d t|D}t|dkrUdddd| Dnd}ddd| D}g}| D]-\}}|ddkr}| d|dqi|dkr| d|qi|dkr| d|qit t|} d} dd | D} dd | D} dd | D} | ddt|dkrd|ndd d| d!| fd"d | Dd#|d$|d%| | d| | d&t|dkrdd|ndd'} | S)(Ncst|tr dt|S|S)Nr4) isinstancetuplermap)sig)_serialize_signaturerrrs z+make_launcher.._serialize_signaturecsJt|trdt|}d|dS|ddkrdS|dvr!dSt|S)Nr4[]rruz PyObject* constexprrrrrrrrval_extracted_typerrrs   z&make_launcher.._extracted_typec sft|trdt|}d|dS|ddkrdS|dvr!dSdd d d d d dddddd t|S)N()rruOrrYdlbhiLBHIK) rrlongrzr|rxrrrrrrr format_ofrrrs*    z make_launcher..format_ofrcsg|]}|qSrr)rrrrrrsz!make_launcher.. iiiKKpOOOOOr4cSi|]\}}||qSrr)rrsrrr z!make_launcher..r, css|] \}}d|VqdS)z&_argNrrrrrrrrsz make_launcher..css.|]\}}|dkrt|d|VqdS)rz argN)rrrrrrs,ruptr_infoz.dev_ptrrz*tma_ptrr_argz c Ss:g|]\}}|ddkrd|d|d|d|d qS)rruzDevicePtrInfo ptr_infoz = getPointer(_argrz); if (!ptr_infoz.valid) return NULL;rrrrrrs  c Ss0g|]\}}|dkrd|d|d|dqS)rzCUtensorMap* tma_ptrz = getTmaDesc(_argz); if (!tma_ptrz) return NULL;rrrrrrs cSs"g|] \}}|dkrd|qS)rz&argrrrrrrs"z&global_scratchaB #include "cuda.h" #include #include #include static inline void gpuAssert(CUresult code, const char *file, int line) { if (code != CUDA_SUCCESS) { const char* prefix = "Triton Error [CUDA]: "; const char* str; cuGetErrorString(code, &str); char err[1024] = {0}; strcat(err, prefix); strcat(err, str); PyGILState_STATE gil_state; gil_state = PyGILState_Ensure(); PyErr_SetString(PyExc_RuntimeError, err); PyGILState_Release(gil_state); } } #define CUDA_CHECK(ans) { gpuAssert((ans), __FILE__, __LINE__); } typedef CUresult (*cuLaunchKernelEx_t)(const CUlaunchConfig* config, CUfunction f, void** kernelParams, void** extra); static cuLaunchKernelEx_t getLaunchKernelExHandle() { // Open the shared library void* handle = dlopen("libcuda.so.1", RTLD_LAZY); if (!handle) { PyErr_SetString(PyExc_RuntimeError, "Failed to open libcuda.so.1"); return NULL; } // Clear any existing error dlerror(); cuLaunchKernelEx_t cuLaunchKernelExHandle = (cuLaunchKernelEx_t)dlsym(handle, "cuLaunchKernelEx"); // Check for errors const char *dlsym_error = dlerror(); if (dlsym_error) { PyErr_SetString(PyExc_RuntimeError, "Failed to retrieve cuLaunchKernelEx from libcuda.so.1"); return NULL; } return cuLaunchKernelExHandle; } static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int launch_cooperative_grid, int clusterDimX, int clusterDimY, int clusterDimZ, int shared_memory, CUstream stream, CUfunction function, CUdeviceptr global_scratchz) { void *params[] = { aO }; if (gridX*gridY*gridZ > 0) { if ((num_ctas == 1) && (0 == launch_cooperative_grid)) { CUDA_CHECK(cuLaunchKernel(function, gridX, gridY, gridZ, 32*num_warps, 1, 1, shared_memory, stream, params, 0)); } else if ((num_ctas == 1) && (0 != launch_cooperative_grid)) { CUlaunchAttribute launchAttr[1]; CUlaunchAttribute coopAttr = { .id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE, .value = 1}; launchAttr[0] = coopAttr; CUlaunchConfig config; config.gridDimX = gridX; config.gridDimY = gridY; config.gridDimZ = gridZ; config.blockDimX = 32 * num_warps; config.blockDimY = 1; config.blockDimZ = 1; config.sharedMemBytes = shared_memory; config.hStream = stream; config.attrs = launchAttr; config.numAttrs = 1; static cuLaunchKernelEx_t cuLaunchKernelExHandle = NULL; if (cuLaunchKernelExHandle == NULL) { cuLaunchKernelExHandle = getLaunchKernelExHandle(); } CUDA_CHECK(cuLaunchKernelExHandle(&config, function, params, 0)); } else { CUlaunchAttribute launchAttr[3]; launchAttr[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; launchAttr[0].value.clusterDim.x = clusterDimX; launchAttr[0].value.clusterDim.y = clusterDimY; launchAttr[0].value.clusterDim.z = clusterDimZ; launchAttr[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; launchAttr[1].value.clusterSchedulingPolicyPreference = CU_CLUSTER_SCHEDULING_POLICY_SPREAD; unsigned numAttrs = 2; if (0 != launch_cooperative_grid) { CUlaunchAttribute coopAttr = { .id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE, .value = 1}; launchAttr[2] = coopAttr; numAttrs = 3; } CUlaunchConfig config; config.gridDimX = gridX * clusterDimX; config.gridDimY = gridY * clusterDimY; config.gridDimZ = gridZ * clusterDimZ; config.blockDimX = 32 * num_warps; config.blockDimY = 1; config.blockDimZ = 1; config.sharedMemBytes = shared_memory; config.hStream = stream; config.attrs = launchAttr; config.numAttrs = numAttrs; static cuLaunchKernelEx_t cuLaunchKernelExHandle = NULL; if (cuLaunchKernelExHandle == NULL) { cuLaunchKernelExHandle = getLaunchKernelExHandle(); } CUDA_CHECK(cuLaunchKernelExHandle(&config, function, params, 0)); } } } typedef struct _DevicePtrInfo { CUdeviceptr dev_ptr; bool valid; } DevicePtrInfo; static inline DevicePtrInfo getPointer(PyObject *obj, int idx) { DevicePtrInfo ptr_info; ptr_info.dev_ptr = 0; ptr_info.valid = true; if (PyLong_Check(obj)) { ptr_info.dev_ptr = PyLong_AsUnsignedLongLong(obj); return ptr_info; } if (obj == Py_None) { // valid nullptr return ptr_info; } PyObject *ptr = PyObject_GetAttrString(obj, "data_ptr"); if(ptr){ PyObject *empty_tuple = PyTuple_New(0); PyObject *ret = PyObject_Call(ptr, empty_tuple, NULL); Py_DECREF(empty_tuple); Py_DECREF(ptr); if (!PyLong_Check(ret)) { PyErr_SetString(PyExc_TypeError, "data_ptr method of Pointer object must return 64-bit int"); ptr_info.valid = false; return ptr_info; } ptr_info.dev_ptr = PyLong_AsUnsignedLongLong(ret); if(!ptr_info.dev_ptr) return ptr_info; uint64_t dev_ptr; int status = cuPointerGetAttribute(&dev_ptr, CU_POINTER_ATTRIBUTE_DEVICE_POINTER, ptr_info.dev_ptr); if (status == CUDA_ERROR_INVALID_VALUE) { PyErr_Format(PyExc_ValueError, "Pointer argument (at %d) cannot be accessed from Triton (cpu tensor?)", idx); ptr_info.valid = false; } else if (status != CUDA_SUCCESS) { CUDA_CHECK(status); // Catch any other cuda API errors ptr_info.valid = false; } ptr_info.dev_ptr = dev_ptr; Py_DECREF(ret); // Thanks ChatGPT! return ptr_info; } PyErr_SetString(PyExc_TypeError, "Pointer argument must be either uint64 or have data_ptr method"); ptr_info.valid = false; return ptr_info; } static inline CUtensorMap* getTmaDesc(PyObject *obj) { if (sizeof(CUtensorMap*) != 8) { PyErr_SetString(PyExc_SystemError, "getTmaDesc() requires 64-bit compilation"); return NULL; } PyObject *method_handle = PyObject_GetAttrString(obj, "tma_desc_cpu_ptr"); if (!method_handle) { PyErr_SetString(PyExc_TypeError, "tma_desc_cpu_ptr() method does not exist"); return NULL; } PyObject *empty_tuple = PyTuple_New(0); if (!empty_tuple) { Py_DECREF(method_handle); PyErr_SetString(PyExc_SystemError, "Internal Python error!"); return NULL; } PyObject *method_ret = PyObject_Call(method_handle, empty_tuple, NULL); Py_DECREF(empty_tuple); Py_DECREF(method_handle); if (!method_ret) { PyErr_SetString(PyExc_SystemError, "Internal Python error!"); return NULL; } if (!PyLong_Check(method_ret)) { PyErr_SetString(PyExc_TypeError, "tma_desc_cpu_ptr() must return 64-bit int"); Py_DECREF(method_ret); return NULL; } uint64_t ptr_as_uint = PyLong_AsUnsignedLongLong(method_ret); Py_DECREF(method_ret); if (!ptr_as_uint) { PyErr_SetString(PyExc_ValueError, "received NULL ptr from tma_desc_cpu_ptr()"); return NULL; } if (ptr_as_uint % 64 != 0) { PyErr_SetString(PyExc_ValueError, "tma_desc_cpu_ptr() must be 64-byte aligned"); return NULL; } return (CUtensorMap*)(ptr_as_uint); } static void ensureCudaContext() { CUcontext pctx; CUDA_CHECK(cuCtxGetCurrent(&pctx)); if (!pctx) { // Ensure device context. CUdevice device; CUDA_CHECK(cuDeviceGet(&device, 0)); CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device)); CUDA_CHECK(cuCtxSetCurrent(pctx)); } } static PyObject* launch(PyObject* self, PyObject* args) { // ensure cuda context is valid before calling any CUDA APIs, e.g. before getPointer calls cuPointerGetAttributes ensureCudaContext(); int gridX, gridY, gridZ; uint64_t _stream; uint64_t _function; int launch_cooperative_grid; PyObject *launch_enter_hook = NULL; PyObject *launch_exit_hook = NULL; PyObject *kernel_metadata = NULL; PyObject *launch_metadata = NULL; PyObject *global_scratch_obj = NULL; cs$g|]\}}|d|dqS)z _arg;rrrrrrs$z if(!PyArg_ParseTuple(args, "a*", &gridX, &gridY, &gridZ, &_stream, &_function, &launch_cooperative_grid, &global_scratch_obj, &kernel_metadata, &launch_metadata, &launch_enter_hook, &launch_exit_hookat)) { return NULL; } int num_warps, num_ctas, shared_memory, clusterDimX, clusterDimY, clusterDimZ; if (!PyArg_ParseTuple(kernel_metadata, "iiiiii", &num_warps, &num_ctas, &shared_memory, &clusterDimX, &clusterDimY, &clusterDimZ)) { PyErr_SetString(PyExc_TypeError, "kernel_metadata must be a tuple"); return NULL; } // extract launch metadata if (launch_enter_hook != Py_None){ PyObject* args = Py_BuildValue("(O)", launch_metadata); PyObject* ret = PyObject_CallObject(launch_enter_hook, args); Py_DECREF(args); if (!ret) return NULL; } CUdeviceptr global_scratch = 0; if (global_scratch_obj != Py_None) { DevicePtrInfo global_scratch_info = getPointer(global_scratch_obj, -1); if (!global_scratch_info.valid) { return NULL; } global_scratch = global_scratch_info.dev_ptr; } // raise exception asap z Py_BEGIN_ALLOW_THREADS; _launch(gridX, gridY, gridZ, num_warps, num_ctas, launch_cooperative_grid, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (CUstream)_stream, (CUfunction)_function, global_scratcha0); Py_END_ALLOW_THREADS; if (PyErr_Occurred()) { return NULL; } if(launch_exit_hook != Py_None){ PyObject* args = Py_BuildValue("(O)", launch_metadata); PyObject* ret = PyObject_CallObject(launch_exit_hook, args); Py_DECREF(args); if (!ret) return NULL; } Py_RETURN_NONE; } static PyMethodDef ModuleMethods[] = { {"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}, {NULL, NULL, 0, NULL} // sentinel }; static struct PyModuleDef ModuleDef = { PyModuleDef_HEAD_INIT, "__triton_launcher", NULL, //documentation -1, //size ModuleMethods }; PyMODINIT_FUNC PyInit___triton_launcher(void) { PyObject *m = PyModule_Create(&ModuleDef); if(m == NULL) { return NULL; } PyModule_AddFunctions(m, ModuleMethods); return m; } ) rvaluesrlistfilterboolr enumeratelenitemsappendrange) constants signature args_formatformat args_list arg_declsinternal_args_listrrparamsnewline ptr_decls tma_declsrQr)rrrr make_launcher~s   ,   ./hil   5rc@seZdZddZddZdS) CudaLaunchercstdrjnt}fddfdd|D}ddjD}t||td}|j|_|j|_|j |_ |j |_ dS)Nrcst|tr jj|fS|Sr-)rr$fn arg_namesindex)x)rQrrsz'CudaLauncher.__init__..csi|] \}}||qSrrridxvalue)arg_idxrrrsz)CudaLauncher.__init__..cSrrrrrrrrr__triton_launcher) rardictrrrr^launchglobal_scratch_sizeglobal_scratch_alignlaunch_cooperative_grid)rorQmetadatarrr]r)rrQrrps    zCudaLauncher.__init__c GsZ|jdkr|||}||j}t||j|} nd} |j||||||j| g|RdSNr)rr _allocatorrrr) rogridXgridYgridZstreamfunctionargs grid_size alloc_sizeglobal_scratchrrr__call__ s   $zCudaLauncher.__call__N)rqrrrsrprrrrrrs rcsXeZdZfddZddZddZddZed d Zd d Z d dZ ddZ Z S) CudaDrivercst|_t|_tdSr-)r_utilsr launcher_clsrbrp)rorerrrpszCudaDriver.__init__cCs6|}||}|dd|d}d}td||S)Nr  r )get_current_deviceget_device_capabilityr)rodevice capability warp_sizerrrget_current_targets   zCudaDriver.get_current_targetcCsddl}|d|S)Nrr )torchrrrorrrrget_active_torch_device"sz"CudaDriver.get_active_torch_devicecCsddl}|jSr)rr rrrrget_device_interface&szCudaDriver.get_device_interfacecCs6zddl}|jo|jjduWStyYdSw)NrF)rr is_availableversionhip ImportError)rrrr is_active*s  zCudaDriver.is_activecCsddlm}|S)Nr)do_bench)triton.testingr)rorrrrget_benchmarker2s zCudaDriver.get_benchmarkercCs&ddl}d}|jt|d|jddS)Nrir )dtyper)remptyint)ror cache_sizerrrget_empty_cache_for_benchmark6sz(CudaDriver.get_empty_cache_for_benchmarkcCs |dSr-)zero_)rorTrrr clear_cache?s zCudaDriver.clear_cache) rqrrrsrprrr staticmethodrrr rrtrrrerrs   r)% functoolsrr@r<r rCpathlibrtriton.runtime.buildrtriton.runtime.cachertriton.runtimertriton.backends.compilerrtriton.backends.driverrrrrealpath__file__rrGr.rH lru_cacher,r/r6r^objectr_rrrrrrrrs>