a jh5c@sddlZddlZddlZddlZddlZddlZddlmZddlmZddl m Z ddl m Z ddl mZddlmZejejeZejedgZejed Zd gZed d Zed dZGdddeZddZddddddZddddddZ dZ!ddZ"GdddZ#e$d d!e%d"DZ&d#e&d$<d$e&d%<d%e&d#<d&d'Z'd(d)Z(Gd*d+d+eZ)Gd,d-d-eZ*dS).N)Path)knobs)compile_module_from_src) _allocation) GPUTarget) GPUDriverincludelibcudacCstjj}r|gStddg}dd|D}dd|D}td}|rj|sjdd| dD}d }|r|d t |7}|d 7}n|d 7}|d 7}t dd|DsJ||S)Nz/sbin/ldconfigz-pcSs g|]}d|vr|dqS) libcuda.so.1)split).0linerK/var/www/auris/lib/python3.9/site-packages/triton/backends/nvidia/driver.py z libcuda_dirs..cSsg|]}tj|qSr)ospathdirname)rlocrrrrrZLD_LIBRARY_PATHcSs&g|]}tjtj|dr|qS)r rrexistsjoin)rdirrrrr r: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(rzlibcuda_dirs..) rZnvidiaZ libcuda_path subprocess check_outputdecode splitlinesrgetenvr strany)Zenv_libcuda_pathZlibsZlocsdirsZenv_ld_library_pathmsgrrr libcuda_dirss    r'cCs tgtSN) libdevice_dirr'rrrr library_dirs,sr*cs$eZdZfddZddZZS) CudaUtilscs"t|dstt|||_|jS)Ninstance)hasattrsuperr+__new__r,)cls __class__rrr/8s zCudaUtils.__new__cCsRtttjtddttt d}|j |_ |j |_ |j |_ |j |_ |j|_dS)Nzdriver.cZ cuda_utilssrcnamer* include_dirs libraries)rrrrrr read_textr*r6r7Z load_binaryZget_device_propertiesZcuOccupancyMaxActiveClustersZset_printf_fifo_sizefill_tma_descriptor)selfmodrrr__init__=szCudaUtils.__init__)__name__ __module__ __qualname__r/r< __classcell__rrr1rr+6s r+cCsH|ddkrdS|drdSddddd d d d d d ddddddd|S)Nr*Z CUdeviceptr tensordescZ CUtensorMapint32_tint8_tint16_tint64_tuint32_tuint8_tuint16_tuint64_tdouble)i1i8Zi16i32i64u1u8u16u32Zu64fp16bf16fp32f32fp64 nvTmaDesc) startswith)tyrrr ty_to_cppQs.  r\rIrGrJ)rTrUrVrWrXZ pack_fp16Z pack_bf16Z pack_fp32Z pack_fp64Z iiiKKppOOOOOcsfdd}fddfddfdd||}d d t|D}d fd d |D}t|}g}|D]}||q~dd t|D}t|dkrdddd|Dnd } g} |D]N\} } | dkrq| tvr| t| d| q| t| d| qd| } g}|D]\} } | ddkrh|d| dnT| tvr|d| dn6| dkr|d| n| dkr>|d| q>t t|}d}dd |D}dd |D}dd |D}d d |D}|d!d"t| dkr:d| nd d#d|d$|fd%d |Dd&|d'| d(||d||d||d)t|dkrdd|nd d*}|S)+Nc sg}d}|D]}t|tr|drڈr0|nd}|d7}td|}|d}|d}|dd}|dur|d|td|D]} |dqn |d t|D]} |d qt|D]} |dqq ||q r|t ksJ|S) NrrBztensordesc<([^[>]*)\[([^]]*)\],rArOrYrN) isinstancer#rZrematchgroupcountappendrangelen) signatureoutputtensordesc_idxsigmetarbdtypeshapendim_)tensordesc_metarr_expand_signature~s,        z(make_launcher.._expand_signaturecs.t|tr |D]}||qn ||dSr()r`tuplere)rkrix)_flatten_signaturerrrus z)make_launcher.._flatten_signaturecsJt|tr&dt|}d|dS|ddkr6dS|dvrBdSt|S)Nr_[]rrAz PyObject* constexprrY)r`rsrmapr\r[val_extracted_typerrr~s   z&make_launcher.._extracted_typec srt|tr&dt|}d|dS|ddkr6dS|dvrBdS|drPdSd d d d d dddddd t|S)N()rrAOrxrBdlbhiLBHIK) rKlongrDrErCrFrHrIrGrJ)r`rsrrzrZr\r{ format_ofrrrs,     z make_launcher..format_ofcSsi|]\}}||qSrrrrsrrr rz!make_launcher..rcsg|] }|qSrr)rr[rrrrrz!make_launcher..cSsi|]\}}||qSrrrrrrrrr, css|]\}}d|VqdS)z&_argNrrrr[rrrrrz make_launcher..ryz argrAZptr_infoz.dev_ptrZ_argZ_storagerYz*tma_ptrz c Ss:g|]2\}}|ddkrd|d|d|d|d qS)rrAzDevicePtrInfo ptr_infoz = getPointer(_argrz); if (!ptr_infoz.valid) return NULL;rrrrrrs c Ss0g|](\}}|dkrd|d|d|dqS)rYzCUtensorMap* tma_ptrz = getTmaDesc(_argz); if (!tma_ptrz) return NULL;rrrrrrsc Ss<g|]4\}}|tvrt|d|dt|d|dqS) _argz _storage = z(_argz);)FLOAT_STORAGE_TYPEFLOAT_PACK_FUNCTIONrrrrrscSs"g|]\}}|dkrd|qS)ryz&argrrrrrrrz&global_scratchaR #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 launch_pdl, int clusterDimX, int clusterDimY, int clusterDimZ, int shared_memory, CUstream stream, CUfunction function, CUdeviceptr global_scratchz) { void *params[] = { a, }; if (gridX*gridY*gridZ > 0) { // 4 attributes that we can currently pass maxmimum CUlaunchAttribute launchAttr[4]; static cuLaunchKernelEx_t cuLaunchKernelExHandle = NULL; if (cuLaunchKernelExHandle == NULL) { cuLaunchKernelExHandle = getLaunchKernelExHandle(); } CUlaunchConfig config; config.gridDimX = gridX; config.gridDimY = gridY; config.gridDimZ = gridZ; if (num_ctas != 1) { config.gridDimX *= clusterDimX; config.gridDimY *= clusterDimY; config.gridDimZ *= clusterDimZ; } config.blockDimX = 32 * num_warps; config.blockDimY = 1; config.blockDimZ = 1; config.sharedMemBytes = shared_memory; config.hStream = stream; config.attrs = launchAttr; int num_attrs = 0; if (launch_pdl != 0) { CUlaunchAttribute pdlAttr = { .id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION, .value = 1}; launchAttr[num_attrs] = pdlAttr; ++num_attrs; } if (launch_cooperative_grid != 0) { CUlaunchAttribute coopAttr = { .id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE, .value = 1}; launchAttr[num_attrs] = coopAttr; ++num_attrs; } if (num_ctas != 1) { CUlaunchAttribute clusterAttr = {}; clusterAttr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; clusterAttr.value.clusterDim.x = clusterDimX; clusterAttr.value.clusterDim.y = clusterDimY; clusterAttr.value.clusterDim.z = clusterDimZ; launchAttr[num_attrs] = clusterAttr; ++num_attrs; CUlaunchAttribute clusterSchedulingAttr = {}; clusterSchedulingAttr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; clusterSchedulingAttr.value.clusterSchedulingPolicyPreference = CU_CLUSTER_SCHEDULING_POLICY_SPREAD; launchAttr[num_attrs] = clusterSchedulingAttr; ++num_attrs; } config.numAttrs = num_attrs; 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 uint16_t pack_fp16(double f) { uint16_t result; // from https://github.com/python/pythoncapi-compat #if 0x030600B1 <= PY_VERSION_HEX && PY_VERSION_HEX <= 0x030B00A1 && !defined(PYPY_VERSION) _PyFloat_Pack2(f, (unsigned char*)&result, 1); #else PyFloat_Pack2(f, (unsigned char*)&result, 1); #endif return result; } static uint16_t pack_bf16(double f) { float f32 = (float)f; uint32_t u32 = *(uint32_t*)&f32; return (uint16_t)(u32 >> 16); } static uint32_t pack_fp32(double f) { float f32 = (float)f; return *(uint32_t*)&f32; } static uint64_t pack_fp64(double f) { return *(uint64_t*)&f; } 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; int launch_pdl; 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)r;rrr}rrrrz if(!PyArg_ParseTuple(args, "a7", &gridX, &gridY, &gridZ, &_stream, &_function, &launch_cooperative_grid, &launch_pdl, &global_scratch_obj, &kernel_metadata, &launch_metadata, &launch_enter_hook, &launch_exit_hooka)) { 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; Py_DECREF(ret); } 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, launch_pdl, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (CUstream)_stream, (CUfunction)_function, global_scratchaC); 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_DECREF(ret); } 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; } ) values enumerater_BASE_ARGS_FORMATrgitemsrrer\rf) constantsrhrqrrZexpand_signatureZ args_formatformatZflat_signaturerkZ args_listZ arg_decl_listrr[Z arg_declsZinternal_args_listparamsnewlineZ ptr_declsZ tma_declsZfloat_storage_declsr4r)r~rurrqr make_launcher|s &     ,       ./$%&(Prc@s eZdZdZddZddZdS)TmaDescKernelParamcCs"ddl}|j|j|jdd|_dS)Nrcpurmdevice)torchempty TMA_DESC_SIZEZuint8descr:rrrrr<VszTmaDescKernelParam.__init__cCs |jSr()rdata_ptrr:rrrtma_desc_cpu_ptr[sz#TmaDescKernelParam.tma_desc_cpu_ptrN)r=r>r?rr<rrrrrrSsrccs|]}||fVqdSr(r)rrrrrr`rr  c Cs|dur(|jg|j|j|j|jS|d}|d}|d}|d}|d}|j}|j}|j} | ddksvJt} | g|| } |rt|}|dd9<tjjj j | |||t |||| | S) Nswizzle elem_size elem_type block_size fp4_paddedr r]r^)basernstridesrrlisttritonZruntimeZdriveractiveutilsr9rTMA_DTYPE_DEVICE_TO_HOST) argmetadatarrrrrrrnrrresultrrrmake_tensordesc_argfs6    rcs.ddlmddlmfdd}|S)Nr)TensorDescriptorcs|dtt}|ttd}d}g}t|D]J\}}t|frprR|nd}|d7}|t||q0||q0r|tksJg||RS)Nrr])rgrrr`extendrre)argsZ meta_argsZraw_kernel_argsrjZ final_argsrrrlZGluonTensorDescriptorrlauncherrqrrinners z%wrap_handle_tensordesc..inner)Ztriton.tools.tensor_descriptorrZ'triton.experimental.gluon.nvidia.hopper)rrqrrrrwrap_handle_tensordescs  rc@seZdZddZddZdS) CudaLaunchercstdrjnt}fddfdd|D}ddjD}t|dd}t|||tdtt t d }t d d | D}t tj|jd |_|rt|j|n|j|_|j|_|j|_|j|_|j|_dS) Nrcst|trjj|fS|Sr()r`r#fn arg_namesindex)rt)r4rrrz'CudaLauncher.__init__..csi|]\}}||qSrrridxvalue)arg_idxrrrrz)CudaLauncher.__init__..cSsi|]\}}||qSrrrrrrrrrqZ__triton_launcherr3css"|]}t|to|dVqdS)rBN)r`r#rZ)rrkrrrrrz(CudaLauncher.__init__..r])r-rdictrrhgetattrrrr*r6r7r$r functoolsreduceoperatormulZ cluster_dimsnum_ctasrlaunchglobal_scratch_sizeglobal_scratch_alignlaunch_cooperative_grid launch_pdl)r:r4rrrhrqr;Zhas_tensor_desc_argr)rr4rr<s(   zCudaLauncher.__init__c Gsd|jdkr8|||}||j|j}t||j|} nd} |j||||||j|j| g|RdSNr)rrrZ _allocatorrrrr) r:ZgridXZgridYZgridZstreamfunctionrZ grid_sizeZ alloc_sizeZglobal_scratchrrr__call__s  zCudaLauncher.__call__N)r=r>r?r<rrrrrrsrcsXeZdZfddZddZddZddZed d Zd d Z d dZ ddZ Z S) CudaDrivercst|_t|_tdSr()r+rrZ launcher_clsr.r<rr1rrr<szCudaDriver.__init__cCs6|}||}|dd|d}d}td||S)Nrrr] r )get_current_deviceZget_device_capabilityr)r:rZ capabilityZ warp_sizerrrget_current_targets  zCudaDriver.get_current_targetcCsddl}|d|S)Nrr )rrrrrrrget_active_torch_devicesz"CudaDriver.get_active_torch_devicecCsddl}|jSr)rr rrrrget_device_interfaceszCudaDriver.get_device_interfacecCs:z ddl}|jo|jjduWSty4YdS0dS)NrF)rr Z is_availableversionZhip ImportError)rrrr is_actives  zCudaDriver.is_activecCsddlm}|S)Nr)do_bench)Ztriton.testingr)r:rrrrget_benchmarkers zCudaDriver.get_benchmarkercCs&ddl}d}|jt|d|jddS)Nrir r)rrint)r:rZ cache_sizerrrget_empty_cache_for_benchmarksz(CudaDriver.get_empty_cache_for_benchmarkcCs |dSr()Zzero_)r:cacherrr clear_cacheszCudaDriver.clear_cache) r=r>r?r<rrr staticmethodrrrrr@rrr1rrs   r)+rrrrrrapathlibrrZtriton.runtime.buildrZtriton.runtimerZtriton.backends.compilerrZtriton.backends.driverrrrrealpath__file__rr6r)r7 lru_cacher'r*objectr+r\rrrrrrrfrrrrrrrrrs^        Z ($