U ,dv @sdZddlZddlZddlZddlZddlmZmZmZm Z m Z m Z m Z ddl Z ddlmZddlmZmZddlmZmZmZddlmZmZeeZdZdZd Zd Z d Z!eZ"eZ#d $Z%e&e%D]\Z'Z(e)ej*ee(e'qd dZ+e ,Z-Gddde.Z/Gddde.Z0dZ1da2ddZ3ddZ4ddZ5dZ6Gddde.Z7dZ8dZ9dZ:d Z;d!Zd$d%Z?d&d'Z@d(d)ZAd*d+ZBd,d-ZCd.d/ZDd0d1ZEd2d3ZFeGd4ZHeGd5ZIeGd6ZJd7ZKeGeKLd8d9ZMeGd:ZNd;dd?d@dAdBdCdDdEdFh ZOeGdGZPeGdHZQeGdIZReGdJZSeGdKZTeGdLZUeGdMZVeGdNZWeGdOZXdPdQdRZYdSdTZZdUdVZ[dWdXZ\dYdZZ]d[d\Z^d]d^Z_d_d`Z`dS)az( This is a direct translation of nvvm.h N)c_void_pc_intPOINTERc_char_pc_size_tbyrefc_char)ir) NvvmErrorNvvmSupportError) get_libdeviceopen_libdevice open_cudalib)cgutilsconfiga NVVM_SUCCESS NVVM_ERROR_OUT_OF_MEMORY NVVM_ERROR_PROGRAM_CREATION_FAILURE NVVM_ERROR_IR_VERSION_MISMATCH NVVM_ERROR_INVALID_INPUT NVVM_ERROR_INVALID_PROGRAM NVVM_ERROR_INVALID_IR NVVM_ERROR_INVALID_OPTION NVVM_ERROR_NO_MODULE_IN_PROGRAM NVVM_ERROR_COMPILATION cCs*z tWntk r YdSXdSdS)z( Return if libNVVM is available FTN)NVVMr rr;/tmp/pip-unpacked-wheel-eu7e0c37/numba/cuda/cudadrv/nvvm.py is_available3s  rc@seZdZdZeeeeefeeefeeefeeee efeeee efeeeeefeeee feeefeeee feeefeeeeeeeeefd Z dZ ddZ ddZ edd Zd d Zd d ZdddZdS)rzProcess-wide singleton. ) nvvmVersionnvvmCreateProgramnvvmDestroyProgramnvvmAddModuleToProgramnvvmLazyAddModuleToProgramnvvmCompileProgramnvvmGetCompiledResultSizenvvmGetCompiledResultnvvmGetProgramLogSizenvvmGetProgramLog nvvmIRVersionNc Cst|jdkrt||_}ztd|_Wn8tk rf}zd|_d}t||W5d}~XYnX|j D]8\}}t |j|}|d|_ |dd|_ t |||qrW5QRX|jS)NZnvvmz;libNVVM cannot be found. Do `conda install cudatoolkit`: %srr ) _nvvm_lock_NVVM__INSTANCEobject__new__rdriverOSErrorr _PROTOTYPESitemsgetattrrestypeargtypessetattr)clsinsteerrmsgnameprotofuncrrrr'ys   z NVVM.__new__cCs4|}|d|_|d|_|d|_|d|_dS)Nrr r)get_ir_version_majorIR_minorIRZ _majorDbgZ _minorDbg)self ir_versionsrrr__init__s    z NVVM.__init__cCs|j|jfdkS)N)r )r9r:r;rrr is_nvvm70szNVVM.is_nvvm70cCs8t}t}|t|t|}||d|j|jfS)NzFailed to get version.)rrr check_errorvalue)r;majorminorerrrrr get_versions  zNVVM.get_versioncCsXt}t}t}t}|t|t|t|t|}||d|j|j|j|jfS)NzFailed to get IR version.)rr#rrArB)r;ZmajorIRZminorIRZmajorDbgZminorDbgrErrrr8s zNVVM.get_ir_versionFcCs2|r.t|t|}|r*t|tdn|dS)Nr )r RESULT_CODE_NAMESprintsysexit)r;errormsgrJexcrrrrAs  zNVVM.check_error)F)__name__ __module__ __qualname____doc__ nvvm_resultrr nvvm_programrrr*r%r'r=propertyr@rFr8rArrrrrBsR   2  rc@sDeZdZddZddZddZddZd d Zd d Zd dZ dS)CompilationUnitcCs4t|_t|_|jt|j}|j|ddS)NzFailed to create CU)rr(rS_handlerrrA)r;rErrrr=szCompilationUnit.__init__cCs*t}|t|j}|j|ddddS)NzFailed to destroy CUT)rJ)rrrrVrA)r;r(rErrr__del__szCompilationUnit.__del__cCs*|j|j|t|d}|j|ddS)z Add a module level NVVM IR to a compilation unit. - The buffer should contain an NVVM module IR either in the bitcode representation (LLVM3.0) or in the text representation. NFailed to add module)r(rrVlenrAr;bufferrErrr add_modules  zCompilationUnit.add_modulecCs*|j|j|t|d}|j|ddS)z Lazily add an NVVM IR module to a compilation unit. The buffer should contain NVVM module IR either in the bitcode representation or in the text representation. NrX)r(rrVrYrArZrrrlazy_add_modules  zCompilationUnit.lazy_add_modulec Ksxg}d|kr |dr |d|ddr6|dd|krR|d|d|drp|d |dd }|D]8}||krxtt||}|d |d d |fqx|rdtt| }t d |t t |dd|D}|j|jt ||}||dt} |j|jt| }||dt| j} |j|j| }||d||_| ddS)aFPerform Compilation The valid compiler options are * - -g (enable generation of debugging information) * - -opt= * - 0 (disable optimizations) * - 3 (default, enable optimizations) * - -arch= * - compute_20 (default) * - compute_30 * - compute_35 * - -ftz= * - 0 (default, preserve denormal values, when performing * single-precision floating-point operations) * - 1 (flush denormal values to zero, when performing * single-precision floating-point operations) * - -prec-sqrt= * - 0 (use a faster approximation for single-precision * floating-point square root) * - 1 (default, use IEEE round-to-nearest mode for * single-precision floating-point square root) * - -prec-div= * - 0 (use a faster approximation for single-precision * floating-point division and reciprocals) * - 1 (default, use IEEE round-to-nearest mode for * single-precision floating-point division and reciprocals) * - -fma= * - 0 (disable FMA contraction) * - 1 (default, enable FMA contraction) * debugz-gZlineinfoFz-generate-line-infooptz-opt=%darchz-arch=%s)ftz prec_sqrtprec_divfmaz-%s=%d_-, zunsupported option {0}cSsg|]}t|dqS)utf8)rencode).0xrrr sz+CompilationUnit.compile..zFailed to compile z&Failed to get size of compiled result.zFailed to get compiled result.N)popappendgetintboolreplacejoinmapreprkeysr formatrrYr(rrV _try_errorrrrrrBr get_loglog) r;optionsoptsZ other_optionskvZoptstrZc_optsrEreslenZptxbufrrrcompiles>#          zCompilationUnit.compilecCs|j|d||fdS)Nz%s %s)r(rAry)r;rErLrrrrx.szCompilationUnit._try_errorcCslt}|j|jt|}|j|d|jdkrht|j}|j|j|}|j|d|j dSdS)Nz#Failed to get compilation log size.r zFailed to get compilation log.rh) rr(r!rVrrArBrr"decode)r;rrEZlogbufrrrry1s   zCompilationUnit.get_logN) rNrOrPr=rWr\r]rrxryrrrrrUs  WrUze-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64cCstrtSzddlm}|}Wnd}YnX|dd|d}d|dd}|dkrfd anF|d krtd an8|d krd an*|d krdan|dkrdand at|tS)Nr)runtime)rr.r z CUDA Toolkit z is unsupported by Numba - z%10.2 is the minimum required version.r) r7) rrrrrrr7rrr>rr>r r>r7rrrr7rr) r) rrrrrrrrrrrr) rrrrrrrrrrrrrr>)rr)rrrrrrrrrrrrr)rr) _supported_ccZnumba.cuda.cudadrv.runtimerrFwarningswarn)rZcudart_versionZctk_verZunsupported_verrrrget_supported_ccsIs0     rcCsvt}|sd}t|t|D]N\}}||kr6|S||kr|dkr\d||}t|q||dSq|dS)z Given a compute capability, return the closest compute capability supported by the CUDA toolkit. :param mycc: Compute capability as a tuple ``(MAJOR, MINOR)`` :return: Closest supported CC as a tuple ``(MAJOR, MINOR)`` zmNo supported GPU compute capabilities found. Please check your cudatoolkit version matches your CUDA version.rz?GPU compute capability %d.%d is not supported(requires >=%d.%d)r )rr enumerate)ZmyccZ supported_ccrLiccrrrfind_closest_arch{s rcCs"tjrtj}n t||f}d|S)z1Matches with the closest architecture option z compute_%d%d)rZ FORCE_CUDA_CCr)rCrDr`rrrget_arch_options rz~Missing libdevice file. Please ensure you have package cudatoolkit >= 10.2 Install package by: conda install cudatoolkit c@s eZdZdZddZddZdS) LibDeviceNcCs0|jdkr$tdkrttt|_|j|_dSN)_cache_r RuntimeErrorMISSING_LIBDEVICE_FILE_MSGrbcr?rrrr=s   zLibDevice.__init__cCs|jSr)rr?rrrrosz LibDevice.get)rNrOrPrr=rorrrrrsrz define internal {T} @___numba_atomic_{T}_cas_hack({T}* %ptr, {T} %cmp, {T} %val) alwaysinline {{ %out = cmpxchg volatile {T}* %ptr, {T} %cmp, {T} %val monotonic ret {T} %out }} z %cas_success = cmpxchg volatile {Ti}* %iptr, {Ti} %old, {Ti} %new monotonic monotonic %cas = extractvalue {{ {Ti}, i1 }} %cas_success, 0 zI %cas = cmpxchg volatile {Ti}* %iptr, {Ti} %old, {Ti} %new monotonic a define internal {T} @___numba_atomic_{T}_{FUNC}({T}* %ptr, {T} %val) alwaysinline {{ entry: %iptr = bitcast {T}* %ptr to {Ti}* %old2 = load volatile {Ti}, {Ti}* %iptr br label %attempt attempt: %old = phi {Ti} [ %old2, %entry ], [ %cas, %attempt ] %dold = bitcast {Ti} %old to {T} %dnew = {OP} {T} %dold, %val %new = bitcast {T} %dnew to {Ti} {CAS} %repeat = icmp ne {Ti} %cas, %old br i1 %repeat, label %attempt, label %done done: %result = bitcast {Ti} %old to {T} ret {T} %result }} a define internal {T} @___numba_atomic_{Tu}_inc({T}* %iptr, {T} %val) alwaysinline {{ entry: %old2 = load volatile {T}, {T}* %iptr br label %attempt attempt: %old = phi {T} [ %old2, %entry ], [ %cas, %attempt ] %bndchk = icmp ult {T} %old, %val %inc = add {T} %old, 1 %new = select i1 %bndchk, {T} %inc, {T} 0 {CAS} %repeat = icmp ne {T} %cas, %old br i1 %repeat, label %attempt, label %done done: ret {T} %old }} a define internal {T} @___numba_atomic_{Tu}_dec({T}* %iptr, {T} %val) alwaysinline {{ entry: %old2 = load volatile {T}, {T}* %iptr br label %attempt attempt: %old = phi {T} [ %old2, %entry ], [ %cas, %attempt ] %dec = add {T} %old, -1 %bndchk = icmp ult {T} %dec, %val %new = select i1 %bndchk, {T} %dec, {T} %val {CAS} %repeat = icmp ne {T} %cas, %old br i1 %repeat, label %attempt, label %done done: ret {T} %old }} a define internal {T} @___numba_atomic_{T}_{NAN}{FUNC}({T}* %ptr, {T} %val) alwaysinline {{ entry: %ptrval = load volatile {T}, {T}* %ptr ; Return early when: ; - For nanmin / nanmax when val is a NaN ; - For min / max when val or ptr is a NaN %early_return = fcmp uno {T} %val, %{PTR_OR_VAL}val br i1 %early_return, label %done, label %lt_check lt_check: %dold = phi {T} [ %ptrval, %entry ], [ %dcas, %attempt ] ; Continue attempts if dold less or greater than val (depending on whether min or max) ; or if dold is NaN (for nanmin / nanmax) %cmp = fcmp {OP} {T} %dold, %val br i1 %cmp, label %attempt, label %done attempt: ; Attempt to swap in the value %old = bitcast {T} %dold to {Ti} %iptr = bitcast {T}* %ptr to {Ti}* %new = bitcast {T} %val to {Ti} {CAS} %dcas = bitcast {Ti} %cas to {T} br label %lt_check done: ret {T} %ptrval }} cCs$tjrtj|dStj|dSdS)NTi)rr@ cas_nvvm70rw cas_nvvm34rrrrir_cas/s rcCs"t||||t|d}tjf|S)N)TrOPFUNCCAS)dictrir_numba_atomic_binary_templaterw)rrrrparamsrrrir_numba_atomic_binary6src Cs&t||||||t|d}tjf|S)N)rrNANr PTR_OR_VALrr)rrir_numba_atomic_minmax_templaterw)rrrrrrrrrrir_numba_atomic_minmax;s  rcCstj||t|dSN)rTur)ir_numba_atomic_inc_templaterwrrrrrrir_numba_atomic_incBsrcCstj||t|dSr)ir_numba_atomic_dec_templaterwrrrrrir_numba_atomic_decFsrcCsD|}t|D](\}}|drd}|t||<q:qd|S)z@ Find the line containing the datalayout and replace it ztarget datalayoutztarget datalayout = "{0}" ) splitlinesr startswithrw data_layoutrs)llvmirlinesrlntmprrr_replace_datalayoutJs rcCsdtdddddfdtdd d d dfd tddd d dfd tdddfdtdddfdtdd dddddfdtdddddddfdtdd dddddfdtdddddddfdtdd dddddfdtdddddddfd tdd dd!dddfd"tdddd!dddfd#g}tjs>|d$tjd d%fd&tjdd%fg7}t|}|D]\}}| ||}qB| d'd(}tjrzt |}nt |}|S))NzIdeclare double @"___numba_atomic_double_add"(double* %".1", double %".2")doubleZi64Zfaddadd)rrrrzEdeclare float @"___numba_atomic_float_sub"(float* %".1", float %".2")floati32ZfsubsubzIdeclare double @"___numba_atomic_double_sub"(double* %".1", double %".2")z=declare i64 @"___numba_atomic_u64_inc"(i64* %".1", i64 %".2")Zu64rz=declare i64 @"___numba_atomic_u64_dec"(i64* %".1", i64 %".2")zEdeclare float @"___numba_atomic_float_max"(float* %".1", float %".2")rznnan oltptrmax)rrrrrrzIdeclare double @"___numba_atomic_double_max"(double* %".1", double %".2")zEdeclare float @"___numba_atomic_float_min"(float* %".1", float %".2")znnan ogtminzIdeclare double @"___numba_atomic_double_min"(double* %".1", double %".2")zHdeclare float @"___numba_atomic_float_nanmax"(float* %".1", float %".2")nanZultzLdeclare double @"___numba_atomic_double_nanmax"(double* %".1", double %".2")zHdeclare float @"___numba_atomic_float_nanmin"(float* %".1", float %".2")ZugtzLdeclare double @"___numba_atomic_double_nanmin"(double* %".1", double %".2"))ZimmargrzMdeclare i32 @"___numba_atomic_i32_cas_hack"(i32* %".1", i32 %".2", i32 %".3"))rzMdeclare i64 @"___numba_atomic_i64_cas_hack"(i64* %".1", i64 %".2", i64 %".3")zllvm.numba_nvvm.atomiczllvm.nvvm.atomic) rrrrrr@ir_numba_cas_hackrwrrrllvm100_to_70_irllvm100_to_34_ir)r replacementsdeclfnrrr llvm_replaceWs          &       rcKsxt|tr|g}|ddr0|dddddt}t}|D]}t|}||dq@| | |j f|S)NZfastmathFT)rardrcrbrh) isinstancestrrmupdaterUrrr\rir]ror)rr|ZcuZ libdevicemodrrr llvm_to_ptxs   rz \!\d+\s*=zmetadata\s*\![{'\"0-9]z\!\d+z,\!{i32 \d, \!\"Debug Info Version\", i32 \d} z\s+z"^attributes #\d+ = \{ ([\w\s]+)\ }Z alwaysinlineZcoldZ inlinehintZminsizeZ noduplicateZnoinlineZnoreturnZnounwindZoptnoneZoptiszeZreadnonereadonlyz"\bgetelementptr\s(?:inbounds )?\(?z=\s*\bload\s(?:\bvolatile\s)?z(call\s[^@]+\))(\s@)z\s*!range\s+!\d+z [,{}()[\]]z \bnonnull\bz"\b(local_unnamed_addr|writeonly)\bz\((.*)\)zspFlags: (.*),Z isDefinitionZ isOptimized)ZDISPFlagDefinitionZDISPFlagOptimizedcCsng}|D]V}|drXt|}|d}ddd|D}||d|}||q d|S)z, Convert LLVM 10.0 IR for LLVM 7.0. attributes #r rcss|]}|dkr|VqdS)Z willreturnNrrjarrr sz#llvm100_to_70_ir..r) rrre_attributes_defmatchgroupsplitrsrrrn)r buflinemattrsrrrrs    rc CsNdd}g}|D],}|dr0|dd}|drRd|krR|dd}t|rdt|kr|d d }|d d }|d }|d|d||dd}}dd}d |t ||f}|drqt |dk rt dd|}|dr@t |}|d} d dd| D} ||d| }d|krt|}|dkrltd|f|} |d| ||| d}d|krt|}|r|} |d| ||| d}d|kr t d|}t d|d}d|kr t t|}d |kr,d|kr,t t|}t d|}||qd! |S)"z, Convert LLVM 10.0 IR for LLVM 3.4. cSsd}d}t||}|dkr,td|fqv|}|d}|dkrR|dkrtqvq|dkrd|d7}q|dkr|d8}q||dS)Nrzfailed parsing leading type: %s,z{[(r z)]}) re_type_toksearchrendrlstrip)sZ par_levelposrtokrrrparse_out_leading_types     z0llvm100_to_34_ir..parse_out_leading_typez!numba.llvm.dbg.cuz !llvm.dbg.cuz%tail call void asm sideeffect "// dbgz !numba.dbgz!dbgNz!{z metadata !{z!"z metadata !"=r cSsd|dS)Nz metadata r)rrrrrfix_metadata_refsz*llvm100_to_34_ir..fix_metadata_refrzsource_filename =cSsdS)Nrrrrrr z"llvm100_to_34_ir..rcss|]}|tkr|VqdSr)supported_attributesrrrrr&sz#llvm100_to_34_ir..zgetelementptr z failed parsing getelementptr: %szload zcall z\1*\2rrz @llvm.memsetZdeclarer)rrrrrre_metadata_defrre_metadata_correct_usagerfindrsre_metadata_refrre_unsupported_keywordsrrrre_getelementptrrrre_loadre_callre_rangerstripre_parenthesized_list_replace_llvm_memset_usage _replace_llvm_memset_declarationre_annotationsrn) r rrrZassigposlhsrhsrrrrrrrrsp       "                rcCsbt|dd}td|d}|s2tdn |d}|dd|d|}d |S) zNReplace `llvm.memset` usage for llvm7+. Used as functor for `re.sub. r rz align (\d+)rz+No alignment attribute found on memset destrzi32 {}rg({})) listrrrer ValueErrorinsertrwrs)rrZ align_attrZalignoutrrrrSs   rcCs4t|dd}|ddd|}d|S)zTReplace `llvm.memset` declaration for llvm7+. Used as functor for `re.sub. r rrrrgr )r rrr rsrw)rrrrrrrcs  rcCsN|j}t|d}ttdd}||||f}t|d}||dS)Nkernel r znvvm.annotations) moduler ZMetaDataStringConstantIntType add_metadatarZget_or_insert_named_metadatar)ZlfuncrZmdstrZmdvalueZmdZnmdrrrset_cuda_kernelns   rcCsftd}tjr0|d|d|d|dg}n|d|d|d|dg}||}|d|dS) zAdd NVVM IR version to modulerr r>rrr7znvvmir.versionN)r rrr@rZadd_named_metadata)rrr<Zmd_verrrradd_ir_versionys   rcCs t|_dSr)r)rrrrfix_data_layoutsr)arQloggingr rIrctypesrrrrrrr threadingZllvmliter rKr r Zlibsr rrZ numba.corerr getLoggerrNloggerZADDRSPACE_GENERICZADDRSPACE_GLOBALZADDRSPACE_SHAREDZADDRSPACE_CONSTANTZADDRSPACE_LOCALrSrRrrGrrr}r/modulesrLockr$r&rrUrrrrrrrrrrrrrrrrrrrrrrrrrrZdebuginfo_patternrrZre_metadata_debuginforrrrrrrrrrZ re_spflagsZ spflagmaprrrrrrrrrrrs$    u 2"   E             k