Ë
    ÎÍ:j5c  ã                   ó’  — d dl Z d dlZd dlZd dlZd dlZd dlZd dlmZ d dlmZ d dl	m
Z
 d dlmZ d dlmZ d dlmZ ej"                  j%                  ej"                  j'                  e«      «      Zej"                  j+                  ed«      gZej"                  j+                  ed	«      Zd
gZ e j2                  «       d„ «       Z e j2                  «       d„ «       Z G d„ de«      Zd„ ZddddddœZddddddœZ dZ!d„ Z" G d„ d«      Z# e$d„  e%d«      D «       «      Z&de&d<   de&d <   d e&d<   d!„ Z'd"„ Z( G d#„ d$e«      Z) G d%„ d&e«      Z*y)'é    N)ÚPath)Úknobs)Úcompile_module_from_src)Ú_allocation)Ú	GPUTarget)Ú	GPUDriverÚincludeÚlibÚcudac            	      óÐ  — t         j                  j                  x} r| gS t        j                  ddg«      j                  «       }|j                  «       D �cg c]  }d|v sŒ|j                  «       d   ‘Œ }}|D �cg c]!  }t        j                  j                  |«      ‘Œ# }}t        j                  d«      }|r^|s\|j                  d«      D �cg c]B  }t        j                  j                  t        j                  j                  |d«      «      sŒA|‘ŒD }}d}|r|dt        |«      z  z  }|d	z  }n
|d
z  }|dz  }t        d„ |D «       «      sJ |«       ‚|S c c}w c c}w c c}w )Nz/sbin/ldconfigz-púlibcuda.so.1éÿÿÿÿÚLD_LIBRARY_PATHú:zlibcuda.so cannot found!
z!Possible files are located at %s.z:Please create a symlink of libcuda.so to any of the files.z<Please make sure GPU is set up and then run "/sbin/ldconfig"z- (requires sudo) to refresh the linker cache.c              3   ó–   K  — | ]A  }t         j                  j                  t         j                  j                  |d «      «      –— ŒC y­w)r   N)ÚosÚpathÚexistsÚjoin)Ú.0r   s     úr/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/triton/backends/nvidia/driver.pyú	<genexpr>zlibcuda_dirs.<locals>.<genexpr>(   s,   è ø€ ÒSÀdŒr�w‰w�~‰~œbŸg™gŸl™l¨4°Ó@×AÑSùs   ‚AA	)r   ÚnvidiaÚlibcuda_pathÚ
subprocessÚcheck_outputÚdecodeÚ
splitlinesÚsplitr   r   ÚdirnameÚgetenvr   r   ÚstrÚany)	Úenv_libcuda_pathÚlibsÚlineÚlocsÚlocÚdirsÚenv_ld_library_pathÚdirÚmsgs	            r   Úlibcuda_dirsr-      sR  € ä Ÿ<™<×4Ñ4Ð4ÐÐ4Ø Ð!Ð!ä×"Ñ"Ð$4°dÐ#;Ó<×CÑCÓE€Dð *.¯©Ó):ÖU ¸nÐPTÒ>TˆD�J‰J‹L˜ÓÐU€DÐUØ,0Ö1 SŒB�G‰G�O‰O˜CÕ Ð1€DÐ1ÜŸ)™)Ð$5Ó6ÐÙ¡4Ø2×8Ñ8¸Ó=Ös˜ÄÇÁÇÁÔPR×PWÑPW×P\ÑP\Ð]`ÐbpÓPqÕAr’ÐsˆÐsØ
&€CÙØÐ2´S¸³YÑ>Ñ>ˆØÐKÑK‰àÐMÑMˆØÐ>Ñ>ˆÜÑSÈdÔSÔSÐXÐUXÓXÐSØ€Kùò VùÚ1ùò ts   Á	EÁ!EÁ<&EÃAE#ÄE#c                  ó$   — t         gt        «       ¢S ©N)Úlibdevice_dirr-   © ó    r   Úlibrary_dirsr3   ,   s   € äÐ+œL›NÐ+Ð+r2   c                   ó$   ‡ — e Zd Zˆ fd„Zd„ Zˆ xZS )Ú	CudaUtilsc                 ód   •— t        | d«      st        t        | �  | «      | _        | j                  S )NÚinstance)ÚhasattrÚsuperr5   Ú__new__r7   )ÚclsÚ	__class__s    €r   r:   zCudaUtils.__new__8   s*   ø€ Ü�s˜JÔ'Ü ¤¨CÑ8¸Ó=ˆCŒLØ�|‰|Ðr2   c                 ó`  — t        t        t        j                  j	                  t
        d«      «      j                  «       dt        «       t        t        ¬«      }|j                  | _
        |j                  | _        |j                  | _        |j                  | _        |j                  | _        y )Nzdriver.cÚ
cuda_utils©ÚsrcÚnamer3   Úinclude_dirsÚ	libraries)r   r   r   r   r   r    Ú	read_textr3   rB   rC   Úload_binaryÚget_device_propertiesÚcuOccupancyMaxActiveClustersÚset_printf_fifo_sizeÚfill_tma_descriptor)ÚselfÚmods     r   Ú__init__zCudaUtils.__init__=   s€   € Ü%Ü”R—W‘W—\‘\¤'¨:Ó6Ó7×AÑAÓCØÜ%›Ü%Üô
ˆð Ÿ?™?ˆÔØ%(×%>Ñ%>ˆÔ"Ø,/×,LÑ,LˆÔ)Ø$'×$<Ñ$<ˆÔ!Ø#&×#:Ñ#:ˆÕ r2   )Ú__name__Ú
__module__Ú__qualname__r:   rL   Ú__classcell__©r<   s   @r   r5   r5   6   s   ø„ ôö
;r2   r5   c                 ó¢   — | d   dk(  ry| j                  d«      ryi dd“dd	“d
d“dd“dd“dd“dd“dd“dd“dd“dd“dd“dd“dd“dd“dd“|    S )Nr   Ú*ÚCUdeviceptrÚ
tensordescÚCUtensorMapÚi1Úint32_tÚi8Úint8_tÚi16Úint16_tÚi32Úi64Úint64_tÚu1Úuint32_tÚu8Úuint8_tÚu16Úuint16_tÚu32Úu64Úuint64_tÚfp16ÚdoubleÚbf16Úfp32Úf32Úfp64Ú	nvTmaDesc)Ú
startswith)Útys    r   Ú	ty_to_cpprr   Q   sï   € Ø	ˆ!�u�‚|ØØ	‡}�}�\Ô"ØðØˆiðàˆhðð 	ˆyðð 	ˆyð	ð
 	ˆyðð 	ˆjðð 	ˆiðð 	ˆzðð 	ˆzðð 	ˆzðð 	�ðð 	�ðð 	�ðð 	ˆxðð 	�ðð  	�]ð!ð" 	ñ#
ð 
r2   re   ra   rh   )ri   rk   rl   rm   rn   Ú	pack_fp16Ú	pack_bf16Ú	pack_fp32Ú	pack_fp64ÚiiiKKppOOOOOc                 óô  ‡‡‡‡— ˆfd„}ˆfd„Šˆfd„Šˆfd„Š ||j                  «       «      }t        |«      D ��ci c]  \  }}||“Œ
 }}}dj                  |j                  «       D �cg c]
  } ‰|«      ‘Œ c}«      }t        |z   }	g }
|j                  «       D ]  } ‰||
«       Œ t        |
«      D ��ci c]  \  }}||“Œ
 }}}t	        |«      dkD  r)ddj                  d„ |j                  «       D «       «      z   nd}g }|j                  «       D ]P  \  }}|d	k(  rŒ|t        v r|j                  t        |   › d
|› �«       Œ2|j                  t        |«      › d
|› �«       ŒR dj                  |«      }g }|j                  «       D ]u  \  }}|d   dk(  r|j                  d|› d�«       Œ$|t        v r|j                  d|› d�«       ŒB|dk(  r|j                  d|› �«       Œ\|d	k7  sŒb|j                  d|› �«       Œw t        t	        |«      «      }d}|j                  «       D ��cg c]  \  }}|d   dk(  rd|› d|› d|› d|› d�	‘Œ }}}|j                  «       D ��cg c]  \  }}|dk(  rd|› d|› d|› d�‘Œ }}}|j                  «       D ��cg c])  \  }}|t        v rt        |   › d|› dt        |   › d|› d�‘Œ+ }}}|j                  «       D ��cg c]  \  }}|d	k7  sŒd|› �‘Œ }}}|j                  d «       d!t	        |«      dkD  rd|z   nd› d"dj                  |«      › d#|j                  |j                  «       D ��cg c]  \  }} ‰|«      › d|› d$�‘Œ c}}«      › d%|	› d&|› d'|j                  |«      › d|j                  |«      › d|j                  |«      › d(t	        |«      dkD  rddj                  |«      z   nd› d)�}|S c c}}w c c}w c c}}w c c}}w c c}}w c c}}w c c}}w c c}}w )*Nc                 ó~  •— g }d}| D �]  }t        |t        «      rû|j                  d«      rê‰
r‰
|   nd }|dz  }t        j                  d|«      }|j                  d«      }|j                  d«      }|j                  d«      dz   }|€9|j                  d|z   «       t        d|z  «      D ]  }	|j                  d«       Œ n|j                  d	«       t        |«      D ]  }	|j                  d
«       Œ t        |«      D ]  }	|j                  d«       Œ �Œ|j                  |«       �Œ" ‰
r|t        ‰
«      k(  sJ ‚|S )Nr   rU   é   ztensordesc<([^[>]*)\[([^]]*)\]é   ú,rS   r^   ro   r]   )
Ú
isinstancer"   rp   ÚreÚmatchÚgroupÚcountÚappendÚrangeÚlen)Ú	signatureÚoutputÚtensordesc_idxÚsigÚmetar   ÚdtypeÚshapeÚndimÚ_Útensordesc_metas             €r   Ú_expand_signaturez(make_launcher.<locals>._expand_signature~   s9  ø€ ØˆØˆð ó 	#ˆCÜ˜#œsÔ#¨¯©°|Ô(DÙ:I� ~Ò6Èt�Ø !Ñ#�äŸ™Ð!CÀSÓI�ØŸ™ A›�ØŸ™ A›�Ø—{‘{ 3Ó'¨!Ñ+�à�<Ø—M‘M #¨¡+Ô.ô # 1 t¡8›_ò -˜ØŸ™ eÕ,ñ-ð —M‘M +Ô.ä˜t›ò )�AØ—M‘M %Õ(ð)ä˜t›ò )�AØ—M‘M %Õ(ò)ð —‘˜cÖ"ð7	#ñ: # n¼¸OÓ8LÒ&LÐLÐLØˆr2   c                 ój   •— t        | t        «      r| D ]  } ‰||«       Œ y |j                  | «       y r/   )r}   Útupler‚   )rˆ   r†   ÚxÚ_flatten_signatures      €r   r“   z)make_launcher.<locals>._flatten_signature¤   s4   ø€ ä�cœ5Ô!Øò .�Ù" 1 fÕ-ñ.ð �M‰M˜#Õr2   c                 ó˜   •— t        | t        «      r!dj                  t        ‰| «      «      }d|› d�S | d   dk(  ry| dv ryt	        | «      S )Nr|   ú[ú]r   rS   z	PyObject*©Ú	constexprro   )r}   r‘   r   Úmaprr   )rq   ÚvalÚ_extracted_types     €r   r›   z&make_launcher.<locals>._extracted_type¬   sT   ø€ Ü�bœ%Ô Ø—(‘(œ3˜°Ó3Ó4ˆCØ�s�e˜1�:ÐØˆa‰5�CŠ<ØØÐ+Ñ+ØÜ˜‹}Ðr2   c                 óØ   •— t        | t        «      r!dj                  t        ‰| «      «      }d|› d�S | d   dk(  ry| dv ry| j	                  d«      ryd	d
dddddddddœ
t        | «         S )NÚ ú(ú)r   rS   ÚOr—   rU   ÚdÚlÚbÚhÚiÚLÚBÚHÚIÚK)
rj   ÚlongrZ   r\   rX   r_   rc   re   ra   rh   )r}   r‘   r   r™   rp   rr   )rq   rš   Ú	format_ofs     €r   r¬   z make_launcher.<locals>.format_of¶   s�   ø€ Ü�bœ%Ô Ø—'‘'œ#˜i¨Ó,Ó-ˆCØ�s�e˜1�:ÐØˆa‰5�CŠ<ØØÐ+Ñ+ØØ�=‰=˜Ô&ØàØØØØØØØØØñ
ô �B‹-ñð 	r2   r�   r   z, c              3   ó,   K  — | ]  \  }}d |› �–— Œ y­w)z&_argNr1   )r   r¥   rq   s      r   r   z make_launcher.<locals>.<genexpr>×   s   è ø€ Ò L±°°B 5¨¨¤Ñ Lùs   ‚r˜   z argrS   Úptr_infoz.dev_ptrÚ_argÚ_storagero   z*tma_ptrz
  zDevicePtrInfo ptr_infoz = getPointer(_argz); if (!ptr_infoz.valid) return NULL;zCUtensorMap* tma_ptrz = getTmaDesc(_argz); if (!tma_ptrz) return NULL;z _argz_storage = z(_argz);z&argz&global_scratchaR  
#include "cuda.h"
#include <stdbool.h>
#include <Python.h>
#include <dlfcn.h>

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;
  ú;z
  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_FORMATr„   ÚitemsÚFLOAT_STORAGE_TYPEr‚   rr   rƒ   ÚFLOAT_PACK_FUNCTION)Ú	constantsr…   rŽ   r�   Úexpand_signaturer¥   Úsrq   Úargs_formatÚformatÚflat_signaturerˆ   Ú	args_listÚarg_decl_listÚ	arg_declsÚinternal_args_listÚparamsÚnewlineÚ	ptr_declsÚ	tma_declsÚfloat_storage_declsr@   r›   r“   r¬   s     `                   @@@r   Úmake_launcherrÇ   |   s	  û€ ô$ôLôôñ. )¨×)9Ñ)9Ó);Ó<ÐÜ"+Ð,<Ó"=×>™$˜!˜Q��A‘Ð>€IÑ>à—'‘'°9×3CÑ3CÓ3EÖF¨R™9 R�=ÒFÓG€KÜ Ñ,€Fà€NØ×ÑÓ!ò 0ˆÙ˜3 Õ/ð0ä"+¨NÓ";×<™$˜!˜Q��A‘Ð<€IÑ<ÜPSÐT]ÓP^ÐabÒPb��t—y‘yÑ L¸)¿/¹/Ó:KÔ LÓLÒLÐhj€Ið €MØ—‘Ó"ò <‰ˆˆ2Ø�ÒØØÔ#Ñ#Ø× Ñ Ô$6°rÑ$:Ð#;¸4À¸sÐ!CÕDà× Ñ ¤I¨b£M ?°$°q°cÐ!:Õ;ð<ð —	‘	˜-Ó(€IØÐØ—‘Ó"ò 	2‰ˆˆ2Øˆa‰5�CŠ<Ø×%Ñ%¨°°°8Ð&<Õ=ØÔ%Ñ%Ø×%Ñ%¨¨Q¨C¨xÐ&8Õ9Ø�;Òà×%Ñ%¨°° nÕ5Ø�;ÓØ×%Ñ%¨¨Q¨C jÕ1ð	2ô ”3�y“>Ó"€Fð €Gð —_‘_Ó&÷áˆAˆrØˆa‰5�CŠ<ð !  Ð#5°a°S¸¸1¸#Ð=MÈaÈSÐPdÒeð€Iñ ð fo×etÑetÓev÷Ù\aÐ\]Ð_aØ�Òð ˜q˜cÐ!3°A°3°oÀaÀSÈÒWð€Iñ ð —_‘_Ó&÷áˆAˆrØÔ#Ñ#ô ˜bÑ!Ð
" %¨ s¨+Ô6IÈ"Ñ6MÐ5NÈeÐTUÐSVÐVXÒYðÐñ ð
 '0§o¡oÓ&7×M™U˜Q ¸2ÀÓ;L��Q�CŠjÐM€FÑMØ
‡M�MÐ#Ô$ð.Fô\ [^ð  _hó  [ið  lmò  [mð  GKð  NWò  GWð  suð  Fvð vØ—y‘y Ó(Ð)ð Q*ðb ‡<�<À	ÇÁÓ@Q×R±u°q¸"‘O BÓ'Ð(¨¨a¨S°Ò2ÓRÓSÐTð U Ø &˜xð (Qð R[ÐP[ð \ð< ‡<�<�	ÓÐð Ø
‡<�<�	ÓÐð Ø
‡<�<Ð#Ó$Ð%ð &Hô qtð  uGó  qHð  KLò  qLð  IMð  PT÷  PYñ  PYð  Zló  Pmò  Imð  RTð  HUð 'UðM	M€Cð\
 €JùóE ?ùâFùó =ùó8ùó
ùóùó
 NùóD Ss5   ½OÁ*OÂ7OÈ)!OÉ O"Ê.O(ËO.Ë#O.ÍO4c                   ó   — e Zd ZdZd„ Zd„ Zy)ÚTmaDescKernelParamé€   c                 óf   — dd l }|j                  | j                  |j                  d¬«      | _        y )Nr   Úcpu©rŠ   Údevice)ÚtorchÚemptyÚTMA_DESC_SIZEÚuint8Údesc©rJ   rÏ   s     r   rL   zTmaDescKernelParam.__init__V  s%   € ÛØ—K‘K × 2Ñ 2¸%¿+¹+Èe�KÓTˆ�	r2   c                 ó6   — | j                   j                  «       S r/   )rÓ   Údata_ptr)rJ   s    r   Útma_desc_cpu_ptrz#TmaDescKernelParam.tma_desc_cpu_ptr[  s   € Ø�y‰y×!Ñ!Ó#Ð#r2   N)rM   rN   rO   rÑ   rL   r×   r1   r2   r   rÉ   rÉ   S  s   „ Ø€MòUó
$r2   rÉ   c              #   ó$   K  — | ]  }||f–— Œ
 y ­wr/   r1   )r   r¥   s     r   r   r   `  s   è ø€ Ò:¨1  A¤Ñ:ùs   ‚é   é
   é   é	   c           
      ó6  — |€=| j                   g| j                  ¢| j                  ¢| j                  ¢| j                  ¢S |d   }|d   }|d   }|d   }|d   }| j                   j                  «       }| j                  }| j                  }	|	d   dk(  sJ ‚t	        «       }
|
g|¢|	¢}|rt        |«      }|dxx   dz  cc<   t        j                  j                  j                  j                  j                  |
j                  «       |||t        |   |||	«       |S )	NÚswizzleÚ	elem_sizeÚ	elem_typeÚ
block_sizeÚ
fp4_paddedr   rz   r{   )Úbaser‹   ÚstridesrÖ   rÉ   ÚlistÚtritonÚruntimeÚdriverÚactiveÚutilsrI   r×   ÚTMA_DTYPE_DEVICE_TO_HOST)ÚargÚmetadatarÞ   rß   rà   rá   râ   rÖ   r‹   rä   rÓ   Úresults               r   Úmake_tensordesc_argrï   f  s'  € ØÐð —‘ÐM˜3Ÿ9™9ÐM s§{¡{ÐM°S·Y±YÐMÀÇÁÐMÐMà�yÑ!€GØ˜Ñ%€IØ˜Ñ%€IØ˜,Ñ'€JØ˜,Ñ'€Jà�x‰x× Ñ Ó"€HØ�I‰I€EØ�k‰k€GØ�2‰;˜!ÒÐÐäÓ€DØÐ%�UÐ%˜WÐ%€FáÜ�U“ˆØˆb‹	�Q‰‹	Ü
‡N�N×Ñ× Ñ ×&Ñ&×:Ñ:Ø×ÑÓØØØÜ  Ñ+ØØØô	ð €Mr2   c                 ó6   ‡ ‡‡‡— ddl mŠ ddlmŠ ˆˆˆ ˆfd„}|S )Nr   )ÚTensorDescriptorc                  óD  •— | d t        t        «       }| t        t        «      d  }d}g }t        |«      D ]N  \  }}t        |‰	‰f«      r*‰r‰|   nd }|dz  }|j	                  t        ||«      «       Œ>|j                  |«       ŒP ‰r|t        ‰«      k(  sJ ‚ ‰
g |¢|¢­Ž S )Nr   rz   )r„   r´   r³   r}   Úextendrï   r‚   )ÚargsÚ	meta_argsÚraw_kernel_argsr‡   Ú
final_argsr¥   rì   r‰   ÚGluonTensorDescriptorrñ   ÚlauncherrŽ   s           €€€€r   Úinnerz%wrap_handle_tensordesc.<locals>.inner’  sÂ   ø€ ØÐ0œ#Ô/Ó0Ð1ˆ	ØœsÔ#4Ó5Ð6Ð7ˆØˆØˆ
Ü Ó0ò 	'‰FˆAˆsÜ˜#Ð 0Ð2GÐHÔIÙ:I� ~Ò6Èt�Ø !Ñ#�Ø×!Ñ!Ô"5°c¸4Ó"@ÕAà×!Ñ! #Õ&ð	'ñ # n¼¸OÓ8LÒ&LÐLÐLÙÐ0˜Ð0 ZÒ0Ð0r2   )Útriton.tools.tensor_descriptorrñ   Ú'triton.experimental.gluon.nvidia.hopper)rù   rŽ   rú   rø   rñ   s   `` @@r   Úwrap_handle_tensordescrý   Ž  s   û€ Ý?Ýa÷1ð €Lr2   c                   ó   — e Zd Zd„ Zd„ Zy)ÚCudaLauncherc                 ó  ‡— t        ‰d«      r‰j                  n	t        «       }ˆfd„}|j                  «       D ��ci c]  \  }} ||«      |“Œ }}}‰j                  j                  «       D ��ci c]  \  }}||“Œ
 }}}t        |dd «      }t        |||«      Št        ‰dt        «       t        t        ¬«      }	t        d„ |j                  «       D «       «      }
t        j                  t        j                   |j"                  d«      | _        |
rt'        |	j(                  |«      n|	j(                  | _        |j*                  | _        |j,                  | _        |j.                  | _        |j0                  | _        y c c}}w c c}}w )Nr¸   c                 ót   •— t        | t        «      r&‰j                  j                  j	                  | «      fS | S r/   )r}   r"   ÚfnÚ	arg_namesÚindex)r’   r@   s    €r   ú<lambda>z'CudaLauncher.__init__.<locals>.<lambda>¨  s-   ø€ ¼ZÈÌ3Ô=O˜SŸV™V×-Ñ-×3Ñ3°AÓ6Ð9€ ÐUV€ r2   rŽ   Ú__triton_launcherr?   c              3   ób   K  — | ]'  }t        |t        «      xr |j                  d «      –— Œ) y­w)rU   N)r}   r"   rp   )r   rˆ   s     r   r   z(CudaLauncher.__init__.<locals>.<genexpr>´  s)   è ø€ Ò!vÐ\_¤*¨S´#Ó"6Ò"W¸3¿>¹>È,Ó;WÓ"WÑ!vùs   ‚-/rz   )r8   r¸   Údictrµ   r…   ÚgetattrrÇ   r   r3   rB   rC   r#   r²   Ú	functoolsÚreduceÚoperatorÚmulÚcluster_dimsÚnum_ctasrý   ÚlaunchÚglobal_scratch_sizeÚglobal_scratch_alignÚlaunch_cooperative_gridÚ
launch_pdl)rJ   r@   rí   r¸   Úarg_idxÚidxÚvaluer…   rŽ   rK   Úhas_tensor_desc_args    `         r   rL   zCudaLauncher.__init__¦  s@  ø€ Ü%,¨S°+Ô%>�C—M’MÄDÃFˆ	ÛVˆØ;D¿?¹?Ó;L×M©Z¨S°%‘W˜S“\ 5Ñ(ÐMˆ	ÑMØ25·-±-×2EÑ2EÓ2G×H¡J C¨�S˜%‘ZÐHˆ	ÑHÜ! (Ð,=¸tÓDˆÜ˜I y°/ÓBˆÜ%ØØ$Ü%›Ü%Üô
ˆô "Ñ!vÐcl×csÑcsÓcuÔ!vÓvÐä!×(Ñ(¬¯©°x×7LÑ7LÈaÓPˆŒÙM`Ô,¨S¯Z©Z¸ÔIÐfi×fpÑfpˆŒØ#+×#?Ñ#?ˆÔ Ø$,×$AÑ$AˆÔ!Ø'/×'GÑ'GˆÔ$Ø"×-Ñ-ˆ�ùó% NùÛHs   ¼E;Á/Fc           
      ó  — | j                   dkD  rF||z  |z  }|| j                  z  | j                   z  }t        j                  || j                  |«      }	nd }	 | j
                  |||||| j                  | j                  |	g|¢­Ž  y ©Nr   )r  r  r   Ú
_allocatorr  r  r  r  )
rJ   ÚgridXÚgridYÚgridZÚstreamÚfunctionrô   Ú	grid_sizeÚ
alloc_sizeÚglobal_scratchs
             r   Ú__call__zCudaLauncher.__call__½  s‹   € Ø×#Ñ# aÒ'Ø ™¨Ñ-ˆIØ" T§]¡]Ñ2°T×5MÑ5MÑMˆJÜ(×3Ñ3°JÀ×@YÑ@YÐ[aÓb‰Nà!ˆNØˆ�‰�E˜5 %¨°¸4×;WÑ;WÐY]×YhÑYhØ"ð	+Ø%)ô	+r2   N)rM   rN   rO   rL   r$  r1   r2   r   rÿ   rÿ   ¤  s   „ ò.ó.+r2   rÿ   c                   óR   ‡ — e Zd Zˆ fd„Zd„ Zd„ Zd„ Zed„ «       Zd„ Z	d„ Z
d„ Zˆ xZS )	Ú
CudaDriverc                 óV   •— t        «       | _        t        | _        t        ‰| �  «        y r/   )r5   rê   rÿ   Úlauncher_clsr9   rL   )rJ   r<   s    €r   rL   zCudaDriver.__init__Ê  s   ø€ Ü“[ˆŒ
Ü(ˆÔÜ‰ÑÕr2   c                 ó~   — | j                  «       }| j                  |«      }|d   dz  |d   z   }d}t        d||«      S )Nr   rÚ   rz   é    r   )Úget_current_deviceÚget_device_capabilityr   )rJ   rÎ   Ú
capabilityÚ	warp_sizes       r   Úget_current_targetzCudaDriver.get_current_targetÏ  sK   € Ø×(Ñ(Ó*ˆØ×/Ñ/°Ó7ˆ
Ø ‘] RÑ'¨*°Q©-Ñ7ˆ
Øˆ	Ü˜ ¨YÓ7Ð7r2   c                 óJ   — dd l }|j                  d| j                  «       «      S )Nr   r   )rÏ   rÎ   r+  rÔ   s     r   Úget_active_torch_devicez"CudaDriver.get_active_torch_deviceÖ  s   € ÛØ�|‰|˜F D×$;Ñ$;Ó$=Ó>Ð>r2   c                 ó"   — dd l }|j                  S r  )rÏ   r   rÔ   s     r   Úget_device_interfacezCudaDriver.get_device_interfaceÚ  s   € ÛØ�z‰zÐr2   c                  ó’   — 	 dd l } | j                  j                  «       xr | j                  j                  d u S # t
        $ r Y yw xY w)Nr   F)rÏ   r   Úis_availableÚversionÚhipÚImportError)rÏ   s    r   Ú	is_activezCudaDriver.is_activeÞ  sC   € ð	ÛØ—:‘:×*Ñ*Ó,ÒL°%·-±-×2CÑ2CÀtÐ2KÐLøÜò 	Ùð	ús   ‚7: º	AÁAc                 ó   — ddl m} |S )Nr   )Údo_bench)Útriton.testingr;  )rJ   r;  s     r   Úget_benchmarkerzCudaDriver.get_benchmarkeræ  s
   € Ý+Øˆr2   c                 ób   — dd l }d}|j                  t        |dz  «      |j                  d¬«      S )Nr   i   é   r   rÍ   )rÏ   rÐ   Úint)rJ   rÏ   Ú
cache_sizes      r   Úget_empty_cache_for_benchmarkz(CudaDriver.get_empty_cache_for_benchmarkê  s.   € Ûð
 'ˆ
Ø�{‰{œ3˜z¨Q™Ó/°u·y±yÈˆ{ÓPÐPr2   c                 ó$   — |j                  «        y r/   )Úzero_)rJ   Úcaches     r   Úclear_cachezCudaDriver.clear_cacheó  s   € Ø�‰�r2   )rM   rN   rO   rL   r/  r1  r3  Ústaticmethodr9  r=  rB  rF  rP   rQ   s   @r   r&  r&  È  s;   ø„ ôò
8ò?òð ñó ðòòQör2   r&  )+r
  r  r   r   ræ   r~   Úpathlibr   r   Útriton.runtime.buildr   Útriton.runtimer   Útriton.backends.compilerr   Útriton.backends.driverr   r   r    ÚrealpathÚ__file__r   rB   r0   rC   Ú	lru_cacher-   r3   Úobjectr5   rr   r¶   r·   r´   rÇ   rÉ   r  rƒ   rë   rï   rý   rÿ   r&  r1   r2   r   ú<module>rQ     sr  ðÛ Û Û 	Û Û Û 	Ý Ý Ý 8Ý &Ý .Ý ,à
�'‰'�/‰/˜"Ÿ'™'×*Ñ*¨8Ó4Ó
5€Ø—‘—‘˜W iÓ0Ð1€Ø—‘—‘˜W eÓ,€ØˆH€	ð €×ÑÓñó ðð. €×ÑÓñ,ó ð,ô;�ô ;ò6
ð4 ØØØØñÐ ð ØØØØñÐ ð #Ð òT÷n	$ñ 	$ñ  Ñ:±°b³	Ô:Ó:Ð Ø Ð ˜Ñ ØÐ ˜Ñ Ø Ð ˜Ñ ò%òPô,!+�6ô !+ôH,�õ ,r2   