a kh @s<ddlmZmZddlZddlZddlZddlZddlZddlZddl Z ddl m Z ddl m Z ddlmZddlmZmZmZmZmZmZmZmZmZmZddlmZddlmZd d lm Z d d l!m"Z"d d l#m$Z$m%Z%m&Z&m'Z'e(de)d  Z*edZ+Gdddej,Z-ddddZ.GdddZ/iZ0gZ1ddZ2d?ddZ3Gdddee+Z4ddZ5d d!Z6d"d#Z7e Gd$d%d%Z8Gd&d'd'e4e+Z9edd(d)d*d+Z:edddddddd,d-d-d.d.d/d/d0d1d2d+Z:d@dddddddd,d3d-d-d.d.d/d/d4d5d6d+Z:Gd7d8d8Z;Gd9d:d:ZZ>dS)A) annotationsdivisionN) defaultdict) dataclass)cached_property) CallableGenericIterableOptionalTypeVarUnionoverloadDictAnyTupleTensorDescriptor) ModuleType)knobs)driver) find_paths_ifget_iterable_pathtype_canonicalisation_dictcanonicalize_dtypez .runtime.jitTcseZdZdZddfdd ZeddZdd Zd d Zd d Z ddZ ddZ ddZ ddZ ddZddZddZddZZS)DependenciesFindera This AST visitor is used to find dependencies of a JITFunction. This can be used to invalidate a JITFunction's hash when its source code -- or that of its dependencies -- changes. This visitor also keeps track of the global variables touched by the JITFunction. When we launch the kernel, we check that these have the same values as they did when we ran this visitor. If not, we raise an error (or otherwise we could recompile). NonereturncsHt||_t|d|_||_||_hd|_ i|_ d|_ dS)Nutf-8> minfloatprintgetattrlistmaxrangelenint isinstanceF) super__init__namehashlibsha256encodehasherglobals nonlocalssupported_python_builtinsused_global_valsvisiting_arg_default_value)selfr-r2r3src __class__@/var/www/auris/lib/python3.9/site-packages/triton/runtime/jit.pyr,)s  zDependenciesFinder.__init__cCs |jSN)r1 hexdigestr7r;r;r<retNszDependenciesFinder.retcCs&t|jrdSt|dd}|tS)NT __module__)inspect isbuiltinfuncr$ startswith TRITON_MODULE)r7noderEmoduler;r;r<_is_triton_builtinRs  z%DependenciesFinder._is_triton_builtinc Cst|tr|j|j@D]\}|\}}|j|\}}|j|\}}||krtd|d|d|jd|jd|d q|j|j|j}|t t |dd7}|j | d dS) NGlobal variable z has value z when compiling z, but inner kernel z has conflicting value z7 from when it was first compiled. This is not allowed.noinlineFr ) r* JITFunctionr5keys RuntimeErrorr-__name__update cache_keystrr$r1r0)r7rEkvar_name_Zv1Zv2Zfunc_keyr;r;r< _update_hashXs &zDependenciesFinder._update_hashcst|jtjur|jS|jjvr&dSfdd}||j\}}|durjst|turt|t st |dds|jj vrt ||fj |jt|f<||S)NcsDj|d}|dur |jfSj|d}|dur@|jfSdS)N)NN)r2getr3)r-valr?r;r< name_lookuprs  z2DependenciesFinder.visit_Name..name_lookupZ__triton_builtin__F)typectxastStoreid local_namesr6rr*rMr$r4copyr5rW)r7rHrZrYZvar_dictr;r?r< visit_Namejs(        zDependenciesFinder.visit_Namecsfdd|jDS)Ncsg|]}|qSr;)visit).0eltr?r;r< z2DependenciesFinder.visit_Tuple..)eltsr7rHr;r?r< visit_TupleszDependenciesFinder.visit_TuplecCs\||j}t|tjr&||j}q |dus>t|ddtkrBdSt||j}|||S)NrPrB) rcvaluer*r] Attributer$rGattrrW)r7rHlhsr@r;r;r<visit_Attributes    z"DependenciesFinder.visit_AttributecCs"dd|jjD|_||dS)NcSsh|] }|jqSr;arg)rdrqr;r;r< rgz7DependenciesFinder.visit_FunctionDef..)argsr` generic_visitrir;r;r<visit_FunctionDefsz$DependenciesFinder.visit_FunctionDefcsnfdd}t|j|j|jr&|jgng|jD]}|q0||j|jdur`|j||j dS)NcsDz6jr Jd_|D]}|dur|qWd_nd_0dS)NTF)r6rc)defaultsexprr?r;r<visit_defaultss z:DependenciesFinder.visit_arguments..visit_defaults) itertoolschain posonlyargsrsvararg kwonlyargsrc kw_defaultskwargrv)r7rHrxrqr;r?r<visit_argumentss (    z"DependenciesFinder.visit_argumentscCs8||}t|tr(|jt|O_n |j|dSr=)rcr*r%r`setadd)r7rHtargetr;r;r<visitAssnTargets  z"DependenciesFinder.visitAssnTargetcCs4t|jdkrtd||jd||dS)Nz2Simultaneous multiple assignment is not supported.r)r(targets TypeErrorrrtrir;r;r< visit_AssignszDependenciesFinder.visit_AssigncCs||j||dSr=rrrtrir;r;r<visit_AnnAssigns z"DependenciesFinder.visit_AnnAssigncCs||j||dSr=rrir;r;r< visit_Fors zDependenciesFinder.visit_For)rPrA __qualname____doc__r,propertyr@rJrWrbrjrorurrrrr __classcell__r;r;r9r<rs % '    rrSrcCsddlmm}t|tr|}|dr^|d}t|}|dsNJd|ddS| dr|dt|ddS|drdt|ddS|drt|dSnJt||j rdt|j St||j r|j }nt|tr|j}nt|}t|d d |S) Nrzconst const**krztl.Z_trB)triton.language.corelanguagecorer*rSstriprF removeprefix _normalize_tyendswithZ pointer_typeZ element_tydtyper-r[rPrrXreplace)tyrr;r;r<rs,         rc@seZdZdZdddddddZedd Zed d d d Zed d ddZeddZ eddZ e ddZ e ddZ dS) KernelParamzBRepresents a parameter (name plus metadata) to a @jit'ed function.r)zinspect.Parameterbool)numparamdo_not_specializedo_not_specialize_on_alignmentcCs||_||_||_||_dSr=)r_paramrr)r7rrrrr;r;r<r, szKernelParam.__init__cCs|jjSr=)rr-r?r;r;r<r-szKernelParam.namerSrcCs(|jjr|jjtjjkrdSt|jjS)NrB)r annotationrC Parameteremptyrr?r;r;r<rszKernelParam.annotationcCsN|j}|dr|dd}n|dr4|dd}|ttvrJ|jSdS)NrrrrrB)rrFrrvalues)r7ar;r;r<annotation_types   zKernelParam.annotation_typecCs d|jvSN constexpr)rr?r;r;r< is_constexpr&szKernelParam.is_constexprcCs |jr dSd|jvp|jdS)NFrr)rrrFr?r;r;r<is_const*szKernelParam.is_constcCs|jjSr=)rdefaultr?r;r;r<r0szKernelParam.defaultcCs|jjtjjkSr=)rrrCrrr?r;r;r< has_default4szKernelParam.has_defaultN)rPrArrr,rr-rrrrrrrr;r;r;r<rs      rcs0ddlmddlmd fdd S) NrrrrFTc s6dur dSttrdSttr|r6d|dnd}dkrJ|rJdSdkrbdkrbd |fSd krzd krzd |fSd |fSnttrdStdrj|f}t|d}|dur|drdndt|d}|t|<|rd|dnd}||fStt rdj fStr,dfStdr.specialize_impl..cstdrt|St|S)N_fields)hasattrr[tuple)valsrpr;r<drgzAcreate_specialize_impl..specialize_impl..cSsg|] }|dqSrr;rr;r;r<rfergcSsg|] }|dqSrr;rr;r;r<rffrgz tensordesc<>,zUnsupported type: %s)r*rr)r"rr dtype2strrXrrMrRrrbaser%Z block_shapeZlayoutrr[) rqrspecialize_valuerkeyZdskresspecZ make_tupleZtysrNinnerZGluonTensorDescriptorrspecialize_extrarrpr<rBsX                  "z/create_specialize_impl..specialize_impl)FTT)rrZ'triton.experimental.gluon.nvidia.hopperr)rr;rr<create_specialize_impl=s  1rFcCs6ttdkrttddtd}|||ddS)Nrc[sdSr=r;)rVkwargsr;r;r<rxrgzmangle_type..)r)r(specialize_impl_cacheappendr)rq specializerr;r;r< mangle_typevs rc@s$eZdZUded<ddddZdS)KernelInterfacerrunrcsfddS)z A JIT function is launched with: fn[grid](*args, **kwargs). Hence JITFunction.__getitem__ returns a callable proxy that memorizes the grid. csj|dd|S)NFgridwarmup)r)rsrrr7r;r<rrgz-KernelInterface.__getitem__..r;)r7rr;rr< __getitem__szKernelInterface.__getitem__N)rPrAr__annotations__rr;r;r;r<r}s rc Csldd|D}ddl}||dd|Dt|dd|Dt||j|d}||}|S)NcSs*i|]"\}}||jjdkr"t|n|qSr)r:rPrSrdrrkr;r;r< rgz1serialize_specialization_data..rcSsg|] }t|qSr;r%rr;r;r<rfrgz1serialize_specialization_data..cSsg|] }t|qSr;rrr;r;r<rfrg)r- signature constant_keys constant_vals attrs_keys attrs_valsoptionsr)itemsjsonrNr%r__dict__dumps) r-r constantsattrsrrrobjZserialized_objr;r;r<serialize_specialization_datas$ rc Cst|jt|ksJg}t|j|D]\}}|jrL|d|dq*|jrVdnd}|jrddnd}|jrrdnd}d|d|d|d|d } |j rt |j t r|j dks|j dd d vrd }|r|d |j d | dn|d |j dq*|| q*dd} dd t t| |jdgdd dd|jDdd |d} dd|jD} t| d<t|j| d<t| | | dS)a2 Equivalent to sig.bind followed by apply_defaults. This generates a native Python function (using exec) which can be memoized on a per-kernel basis to avoid having to run these expensive functions -- which constitute much of the kernel launch overhead -- every time we run the kernel. z("constexpr", )TrueFalsezspecialize_impl(, rNr)fpZbfFz("z",) + z[1:]z", None)cSs0|djtjjur|dS|dd|dS)Nrrz =default_rrCrr)rr;r;r<rrgz0create_function_from_signature..z def dynamic_func(z **optionsz): params = {cSsg|]}d|d|qS)'z': r;)rdr-r;r;r<rfrgz2create_function_from_signature..z} specialization = [rz-] return params, specialization, options cSs,i|]$\}}|jtjjurd||jqS)Zdefault_r)rdr-rr;r;r<rsz2create_function_from_signature..rMrZ dynamic_func)r( parametersziprNrrrrrrr*rSjoinr%maprrMrZget_arg_specializationexec) sigZkparamsbackendspecializationr-Zkprrrr@rqZ func_bodyZfunc_namespacer;r;r<create_function_from_signatures@  rcCs|jd|jS)N.)rArfnr;r;r< get_full_namesrc@s&eZdZUded<ded<ded<dS)JitFunctionInforrIrSr-rMZ jit_functionN)rPrArrr;r;r;r<rs rcseZdZddZddddZddZd d Zd d Zd dZd&ddZ ddZ e ddZ e ddZ ddZddZddZddZfd d!Zfd"d#Zd$d%ZZS)'rMcCsdS)NFr;r?r;r;r<is_gluonszJITFunction.is_gluonz bool | Nonerc Cs|sdS|jj} |jj} dddt|j|dD} | d|jd|jd|jd|j d |j d | d } t |j} t | |||d ||}||||j|j|j|j |j |j |||d }||| t| | |d|i||ddS)NrcSs g|]\}}|jd|qS)z: r-)rdrrr;r;r<rfrgz*JITFunction._call_hook..rz [num_warps=z , num_ctas=z , num_stages=z, enable_fp_fusion=z, launch_cooperative_grid=](rr) rdevicer num_warpsnum_ctas num_stagesenable_fp_fusionlaunch_cooperative_grid extern_libsconfigsspecialization_data is_warmuprF)rreprrcompileZis_manual_warmupZalready_compiled)rrrArrparamsrrrrrrrr r)r7hookrrrrrr r r-rIZ arg_reprsr Z full_namer rr;r;r< _call_hooks:  8   zJITFunction._call_hookcCst|s J|j|dS)z Add a hook that will be executed prior to the execution of run function with args and kwargs passed into the kernel N)callable pre_run_hooksr)r7rr;r;r<add_pre_run_hooks zJITFunction.add_pre_run_hookcCsXddlm}m}m}m}tj}||}||_||_||_t|j |j |}i|||fS)z1 Precompute as much as possible. r)CompiledKernelr ASTSource make_backend) compilerrrrrractiveZget_current_targetrrr)r7rrrrrrbinderr;r;r< create_binders zJITFunction.create_binderc! s|d|jptjj|d<tj}tj|}|jD]}||i|q6|j |\}} } | |i|\} } t | t | } || d}|dur |} dd|j D}dd| D}ddt ||D}d|vsJdd |vsJd d |vsJd |D](}|| jvr||vrtd |qt|dd}fdd|D}dd| Dtdd}fdd|D}|tjj| |||| |g|rdS|||||}|j|| | jd}||| <|tjj| |||| |g|t}|jD]B\\}}\}}|||}|krtd|d|d|q|s|dusVJt|rh|}t|}|d}|dkr|dnd}|dkr|dnd}|j||gR} |j|||||j|j | tjj!tjj"g R|S)NdebugcSsg|] }|jqSr;rrr;r;r<rf<rgz#JITFunction.run..cSsg|] }|dqSrr;rr;r;r<rf=rgcSsi|]\}}||qSr;r;)rdrTvr;r;r<r>rgz#JITFunction.run..Z device_typez=device_type option is deprecated; current target will be usedrz8device option is deprecated; current device will be usedstreamz8stream option is deprecated; current stream will be usedz2Keyword argument %s was specified but unrecognisedcSs|dkSrr;)rVrYr;r;r<rGrgz!JITFunction.run..cs i|]}|tt|qSr;)rr%r)rdpath) bound_argsr;r<rHrgcSsg|] }|dqSrr;rr;r;r<rfJrgcSs t|tSr=)r*rS)rVrr;r;r<rKrgcsi|]}|t|qSr;)Z parse_attrr)rdrT)attrvalsrr;r<rLrg)rrrKz1 has changed since we compiled this kernel, from z to rrr)#rXrrruntimerrget_current_deviceZget_current_streamr device_cachesrSZ parse_optionsrrrKeyErrorrrZjit_cache_hookrrZjit_post_compile_hookobjectr5rrOrr(launch_metadatarrfunctionZpacked_metadataZlaunch_enter_hookZlaunch_exit_hook)!r7rrrsrrrrZ kernel_cacherrrrrkernelZsigkeysZsigvalsrrT constexprsrr8Z not_presentr-rVrYZ globals_dictZnewValZ grid_sizeZgrid_0Zgrid_1Zgrid_2r'r;)r!rr r<r#sp        zJITFunction.runcCs|jdur|jS||Sr=)_repr_fn_name)r7rVr;r;r<r mszJITFunction.reprNc Csz|r|ng}|r|ng}||_|j|_||_t||_||_||_t|d|_ ||_ t ||_ ||_ g|_t|jjD]B\} } | |vp| j|v} | |vp| j|v} |jt| | | | qtt|} | td| tjd} || t|j|_d|_ i|_!d|_"||_#||_$dd|jD|_%dd|jD|_&g|_'|j(|_(|j)|_)|j*|_*|j+|_+|j|_dS)Nrz^def\s+\w+\s*\(cSsg|] }|jqSr;rrdpr;r;r<rfrgz(JITFunction.__init__..cSsg|]}|jr|jqSr;)rrr-r;r;r<rfrg),rrArIversionrCrrrgetsourcelinesstarting_line_numberr+rr,r'r enumeraterrr-rrtextwrapdedent getsourceresearch MULTILINEstart_unsafe_update_srcrrr$hashr5r)rrL arg_namesr*rrrPr __globals__)r7rr/rrrrLr r'irZdnsZdns_oar8r;r;r<r,psD       zJITFunction.__init__cCs|jt|jjBSr=)r=rCgetclosurevarsrr3r?r;r;r<get_capture_scopeszJITFunction.get_capture_scopecCsh|jdurbt|jj}t|j|j||jd}| | |j t |j |_tt|j|_|jS)N)r-r2r3r8)r;rCr?rr3rr,r=r8rcparser@rSr1dictsortedr5r)r7r3Zdependencies_finderr;r;r<rRs  zJITFunction.cache_keycCsddlm}|S)Nrr)rr)r7rr;r;r<r[s zJITFunction.typecOs|jttj||dd|S)NTr)rr MockTensor wrap_dtype)r7rrsrr;r;r<rszJITFunction.warmupcsddlm}m}ddl}ddlmtj}| |}|d|j kr`t d|dd|j t t |d}|d}fd d t||D} t t |d } |d } tt| | } t|d } ||| | | }dd |dD}|d}||d|}||j|d|<|S)Nr)rrrr-zSpecialization data is for z but trying to preload for rrcs,i|]$\}}|j|r$|n|qSr;)rZis_dtypertlr;r<rsz'JITFunction.preload..rrrcSs(i|] \}}|t|tr t|n|qSr;)r*r%rrr;r;r<rsrr)rrrrZtriton.languagerrrr#loadsr,rOrrrrBrr$)r7r rrrrZdeserialized_objrrrrrrrr8rrr)r;rFr<preloads4      zJITFunction.preloadcCsHt|j}t|tjsJt|jdks.Jt|jdtjsDJ|S)Nrr)r]rAr8r*Moduler(body FunctionDef)r7treer;r;r<rAs  zJITFunction.parsecOs tddS)Nz:Cannot call @triton.jit'd outside of the scope of a kernel)rO)r7rsrr;r;r<__call__szJITFunction.__call__cs.|dkrtd|dtt|||dS)Nr8zCannot set attribute 'zX' directly. Use '_unsafe_update_src()' and manually clear `.hash` of all callersinstead.)AttributeErrorr+rM __setattr__)r7r-rkr9r;r<rPszJITFunction.__setattr__csd|_td|dS)z The only method allowed to modify src. Bypasses the __setattr__ restriction by calling super().__setattr__ directly. Nr8)r;r+rP)r7Znew_srcr9r;r<r:szJITFunction._unsafe_update_srccCsd|jd|jjdS)Nz JITFunction(:r)rIrrr?r;r;r<__repr__szJITFunction.__repr__)NNNNNNN)rPrArrrrrrr r,r@rrRr[rrIrArNrPr:rRrr;r;r9r<rMs*. J >     rMJITFunction[T]rrcCsdSr=r;rr;r;r<jitsrUr/r r'rrrrLzOptional[Callable]zOptional[Iterable[int | str]]zOptional[bool]zCallable[[T], JITFunction[T]])r r'rrrrLrcCsdSr=r;rVr;r;r<rU s z Optional[T]z4Union[JITFunction[T], Callable[[T], JITFunction[T]]])rr r'rrrrLrc s8dddfdd }|dur0||S|SdS)a< Decorator for JIT-compiling a function using the Triton compiler. :note: When a jit'd function is called, arguments are implicitly converted to pointers if they have a :code:`.data_ptr()` method and a `.dtype` attribute. :note: This function will be compiled and run on the GPU. It will only have access to: * python primitives, * builtins within the triton package, * arguments to this function, * other jit'd functions :param fn: the function to be jit-compiled :type fn: Callable rrSrTc sTt|s Jtjjr8ddlm}||dSt|dSdS)Nr)InterpretedFunction)r/rrrrLr r')rrr"Z interpret interpreterrWrM)rrWrrrr'rLr r/r;r< decorator8s"  zjit..decoratorNr;) rr/r r'rrrrLrZr;rYr<rUs c@s<eZdZdZeddZddZeddZedd Zd S) rDzr Can be used in place of real tensors when calling: kernel.warmup(MockTensor(torch.float32), ...) cCs"|jjdkr|jdkrt|S|S)NrZtorch)r:rPrArDrpr;r;r<rE]szMockTensor.wrap_dtypecCs ||_dSr=r)r7rr;r;r<r,cszMockTensor.__init__cCsdSNrr;r;r;r;r<rfszMockTensor.data_ptrcCsdSr[r;r;r;r;r< ptr_rangejszMockTensor.ptr_rangeN) rPrArr staticmethodrEr,rr\r;r;r;r<rDWs  rDc@sbeZdZddZddZddZddd d Zd d Zd dZddZ ddZ ddZ ddZ dS) TensorWrappercCs*||_||_|j|_|j|_|jj|_dSr=)rrdatarshape)r7rrr;r;r<r,qs zTensorWrapper.__init__cCs |jSr=)rrr?r;r;r<rxszTensorWrapper.data_ptrcGs |jj|Sr=)rstride)r7rsr;r;r<ra{szTensorWrapper.striderSrcCsd|jd|jdS)NzTensorWrapper[rr)rrr?r;r;r<__str__~szTensorWrapper.__str__cCs |jSr=)r element_sizer?r;r;r<rcszTensorWrapper.element_sizecCst|j|jSr=)r^rcpurr?r;r;r<rdszTensorWrapper.cpucCs|j|jdSr=)rcopy_)r7otherr;r;r<reszTensorWrapper.copy_cCst|j|jSr=)r^rclonerr?r;r;r<rgszTensorWrapper.clonecCst|j||jSr=)r^rtor)r7rr;r;r<rhszTensorWrapper.tocCst|j||jSr=)r^r new_emptyr)r7sizesr;r;r<riszTensorWrapper.new_emptyN) rPrArr,rrarbrcrdrergrhrir;r;r;r<r^osr^cCsVt|tr*||jjkr|jSt|j|Sn(t|dr>t||Stdt|ddS)NrzCannot reinterpret a r)r*r^rrrrr[)rrr;r;r< reinterprets    rkcCsd|}t|ts|j}q|jjj}t|j\}}t|D]"\}}| dr8||7}q\q8||fS)Nzdef ) r*rMr__code__ co_filenamerCr0r2rrF)rZbase_fn file_namelinesZ begin_lineidxliner;r;r<get_jit_fn_file_lines  rr)F)N)? __future__rrr]rar.rCryr6r3 collectionsrZ dataclassesr functoolsrtypingrrr r r r r rrrZtriton.tools.tensor_descriptorrtypesrrBrZruntime.driverr_utilsrrrrrPr(rGr NodeVisitorrrrrrrrrrrrrrMrUrDr^rkrrr;r;r;r<sx   0    Q29   :/  <%