L iWddlmZddlZddlZddlZddlmZmZmZddl m Z ddl m Z ddl Z ddlmcmZddlmZddlmZddlmZddlmZdd lmZmZdd lmZmZmZm Z m!Z!m"Z"dd l#m$Z$dd l%m&Z&m'Z'm(Z(m)Z)m*Z*dd l+m,Z,e rddl-m.Z.edgdZ/ ddZ0 d ddZ1ddZ2 d ddZ3 d ddZ4 d d dZ5e6dk(re2yy)!) annotationsN)Counter defaultdict namedtuple)Path) TYPE_CHECKING)DispatcherSignature) CodeTemplate)native_function_manager)get_grouped_native_functionsparse_native_yaml) BackendIndexBackendMetadata DispatchKeyNativeFunctionNativeFunctionsGroup OperatorName)SelectiveBuilder) concatMapcontext FileManagerNamespaceHelperTarget) YamlLoader)SequenceParsedExternalYaml) backend_key autograd_key class_name cpp_namespacebackend_indicesc td|Dcic]}|jj|c}t|5}t j |t }dddttsJgd}|jddJd|jdd}|jddJd|jd d }t|ts Jd ||jd d }t|ts Jd ||jdg} | g} t| tsJd| dt| d|jdg} | g} t| tsJd| dt| dt| } |jdg} t| ts Jd| |jdg} | j| |jdi|jdit|j!dk(s;J|ddj#|j!ddj#| d%fd }d}t| dkDrDt%fd5t'j(}ddd|| | |||}||vsJ|||<d}t| dkDrGt%fd 5t'j(d!}ddd|| | |||}||vsJ|||<|D]X}t|t*rS|gn$||j-|fDcgc]}||c}}|gn$||j-|fDcgc]}||c}}n|gn?|j/Dcgc]}||j-|c}Dcgc]}||c}}|gn?|j/Dcgc]}||j-|c}Dcgc]}||c}}|Dcgc]}|| }}|Dcgc]}|| }}t|dk(r#t|dk(r3Jd"|dj0d#|dj0d$t3||||Scc}w#1swY?xYw#1swYxYw#1swYxYwcc}wcc}wcc}wcc}wcc}wcc}wcc}wcc}w)&NcZt|tr|gSt|jSN) isinstancerlist functionsfs `/mnt/ssd/data/python-lab/Trading/venv/lib/python3.12/site-packages/torchgen/gen_backend_stubs.pyz$parse_backend_yaml..3s Z>:qcQ[[]@S)Loader) backendrr extra_headers supportedautograd full_codegen non_nativeir_gensymintr.z&You must provide a value for "backend"rr z,You must provide a value for "cpp_namespace"use_out_as_primaryFzHYou must provide either True or False for use_out_as_primary. Provided: device_guardzBYou must provide either True or False for device_guard. Provided: r0z,expected "supported" to be a list, but got: z (of type )r5z)expected "symint" to be a list, but got: r1z+expected "autograd" to be a list, but got: r2r3r4rz contains unexpected keys: z, z). Only the following keys are supported: ci}|D]c}tj|}| vs Jd|tj |j}||vr|dz }t |d } | ||<et ||d||S)Nz Found an invalid operator name: _symintF)kernel structuredr T) dispatch_keyr6externalr7index)rparse dispatchernamefuncrr) backend_ops symint_opsr=r6use_device_guardmetadataopop_name kernel_namemr native_functions_maps r*create_backend_indexz0parse_backend_yaml..create_backend_indexs9; "B"((,G22 27)< 2%//*>w*G*L*LMKZy( "uMA!"HW  "%1)   r,cddS)NzFThe provided value for "backend" must be a valid DispatchKey, but got .r.sr*r+z$parse_backend_yaml..s\]d\eefgr,)r6rFcddS)NzThe "autograd" key was specified, which indicates that you would like to override the behavior of autograd for some operators on your backend. However "Autogradz" is not a valid DispatchKey.rPrQsr*r+z$parse_backend_yaml..sOOViWtvr,AutogradzCurrently, all variants of an op must either be registered to a backend key, or to a backend's autograd key. They cannot be mix and matched. If this is something you need, feel free to create an issue! z" is listed under "supported", but z is listed under "autograd".) rDz list[str]rEzset[str]r=rr6boolrFrTreturnr)rrCrBopenyamlloadrr%dictpoprTr&typesetextendlenkeysjoinrrr@r get_kernelr'r;r)backend_yaml_pathgrouped_native_functionsr!r) yaml_values valid_keysrr6rFr0r5 symint_setsupported_autogradr2rMr backend_idxr autograd_idxgrKforward_kernelsbackward_kernelsr.r rLs @@@r*parse_backend_yamlrm+s  S $ @   Q@  6Aii*5 6 k4 (( ( Jooi.G  H HH t4JOOOT:M  $T&TT $%)=uE ($ / RSeRfg /#~u= & - LM]L^_ - R0I i & 6ykDQZOK\\]^ &__Xr *F ~ fd # 3I;jiHYYZ[ #VJ$R8 ($ / 56H5IJ / ??>26L \"OOL"%OOHb! {! "a '  8;CSCSCU9V8WX226))J2G1H J '    " !      >'+K 9~  g  5&++G4K 5 +   1-  /111'2 $'+L "  v  C',,xy-ABL  C ,   1-  ?222(4 % %4  a (&.k:EEaHI}  '.l;FFqIJ} & "#( 4??B}   ' "#( 5@@C}  '6G1GG'7I!1=AII?#q(C0@,AQ,F =>Nq>Q>X>X=YYu w Fa4 l \:}o O@66P 5 5" C C*HIsdRRR%-R2 R? S%S  S)S SS#S-S"5S"R"%R/2R<c  t|5}|j}ddd|g}||jg|gn||jgz} t t td| } |Dcgc]B}|jj| jvr|jj|vr|D} }tt } | D]-}| | |jjj|/d|d}ttj|Dcgc]\}}|j!ds|c}}}d}| j#D]Y\}}t%|}||}||k7sdd}dj'|Dcgc] }|| c}}|d|d |d |d |d |d z }[|dk(sJ|y#1swYxYw#t$r} td|| d} ~ wwxYwcc}wcc}}wcc}w)Nz2Unable to read from the specified impl_path file: cl|jDcgc]\}}||jfc}}Scc}}wr$)itemsr;)r?rIrGs r*r+z*error_on_missing_kernels..s/HM 3D7HWhoo.s0z(.*)z::\s*([\w\d]*)\(:ct|5tj|jj cdddS#1swYyxYwr$)r r from_schemarCdeclr(s r* create_declz-error_on_missing_kernels..create_declCs>,Q/J.::166BGGIJJJs -AA  z$ is missing a kernel definition for z . We found z( kernel(s) with that name, but expected zH kernel(s). The expected function schemas for the missing operator are: z )r)rrUstr)rVreadOSErrorAssertionErrorr?rYr&rrCrBr_rappendrrefindallendswithrpr^r`)native_functionsr!rrrkernel_defn_file_pathr2r) backend_defnseindicesexpected_backend_op_namesexpected_backend_native_funcs#expected_backend_kernel_name_countsnative_fkernel_defn_regexxy!actual_backend_kernel_name_countsmissing_kernels_err_msg expected_namefuncsexpected_overload_countactual_overload_countrvexpected_schemas_strs r*error_on_missing_kernelsrsw ' ( %AFFHM % {+112")F)L)L(MG :>     :"; 66;;388: : FFKK| + ;!; LW L'2+ %hmm&8&8 9 &  |+;<(/**%6 F A::c?  )%! C I I K u"%e* A- P "&; ; J$(99e-Lk!n-L#M #, 0{K`Jab% &'(  # #b (A*AA (Q % % @AV@W X  .;. ".Ms@ G GG 2AG- G2 G8 G G G*G%%G*cxtjd}|jdddd|jdd d d |jd d tdd |jddtdd|j }t |j|j|j|jy)NzGenerate backend stub files) descriptionz-sz --source-yamlz --source_yamlzApath to source yaml file containing operator external definitions)helpz-oz --output-dirz --output_dirzoutput directoryz --dry-runz --dry_runF)r[defaultrz --impl-pathz --impl_pathz9path to the source C++ file containing kernel definitions) argparseArgumentParser add_argumentrTrx parse_argsrun source_yaml output_dirdry_run impl_path)parseroptionss r*mainrQs  $ $1N OF   P   nnCUV [tUAS   H !GW//'BSBSTr,c  Jd tttfd| tttfd|t| |j dd  fdy)NzAAutogenerated file by gen_backend_stubs.py. Do not edit directly!c6tj|Sr$dest#compute_native_function_declaration)r)backend_dispatch_keyr!s r*r+z4gen_dispatchkey_nativefunc_headers..~s$BB';<r,c>gStj|Sr$r)r)autograd_dispatch_keyr!s r*r+z4gen_dispatchkey_nativefunc_headers..s.(0=='<=r,NativeFunctions.hzDispatchKeyNativeFunctions.hcFjjzdS)N)generated_commentnamespace_prologuernamespace_epiloguedispatch_declarations BackendNamer)prologueepilogue)autograd_declarationsbackend_declarationsr backend_namerr ns_helpersr*r+z4gen_dispatchkey_nativefunc_headers..s1!2"+"4"4$"+"4"4%9NrP.0hs r* z/gen_dispatcher_registrations..0Tq:aS1B0Tc3(K|] }d|d yw)z #include ""NrPrs r*rz/gen_dispatcher_registrations..rrFTrocmr5class_method_nameskip_dispatcher_op_registrationat) namespace_strrrzOTORCH_LIBRARY_IMPL(aten, $dispatch_key, m) { $dispatch_registrations_body })r=dispatch_registrations_bodyzTORCH_API void Register${backend_name}${dispatch_key}NativeFunctions(); TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() { static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key); $dispatch_registrations_body })rr=rRegisterz.cppzRegisterDispatchKey.cppc d sdndjtj dtjj d fdj dS)Nrrz#include F)per_operator_headersrzRegisterDispatchDefinitions.inicjjjdtt t j tjddddS)NrrFTr) ns_prologue ns_epilogue"static_init_dispatch_registrationsdeferred_dispatch_registrationsdispatch_namespacedispatch_namespaced_definitionsdispatch_anonymous_definitions) rrlowerr&rrRegisterDispatchKeyrANONYMOUS_DEFINITION) backend_indexrrr=rcrselectorrsr*r+z@gen_dispatcher_registrations....sq#,#5#5#,#5#5:\7V*6*<*<*>796:! 44 - & ; ; (%*'+5?L@E5  7r,)extra_cuda_headersexternal_backend_headers ops_headersrrdispatch_headersdispatch_helpersdispatch_definitions)rrgen_registration_headersgen_registration_helperssubstitute_with_templatesplit) rrrr=external_backend_headers_strrrcnewlinerrrrs r*r+z.gen_dispatcher_registrations..sw"$(D'9'"."4"4"6 $ = =4Hu!!% = =m L$&$?$?1%0eGnI% r,) r`r&rrrr REGISTRATIONrr substituter)rrrr!rcrr=r build_in_treerreager_registrationheadersrstatic_templatedeferred_templaterrrrrrs` ` ` `` ` @@@@@@r*gen_dispatcher_registrationsrsS ,a,-->?G'+yy0TG0T'T$'+yy0TG0T'T$  !! !#L1M"&  $ $##%/L05  %  #Gd3I&(#)+&&   .=-G-G%(C.H. * )  +<*F*F%%(C+G+ '  <.%!% % )r,c ttjjj}tj j |ddfd }||}tj j |d}tj j |d}t||} | j| j} } t| } t|| | } | j}| j}| j}| j}| j} t!j"}|y|| |j%}|J|t'| | ||||t)|||| | ||||gn||gD]}t+|||| | |||y)Nzaten/src/ATen/templatesc t|S)N) install_dir template_dirr)r)rrrs r*make_file_managerzrun..make_file_managers#,  r,z*aten/src/ATen/native/native_functions.yamlzaten/src/ATen/native/tags.yaml)rrxrUr)r__file__absoluteparentospathr`r rr!r rmrrr rrget_nop_selectornative_function_class_namerrr)rrrr pytorch_rootrrnative_yaml_pathtags_yaml_path parsed_yamlrr!rcparsed_backend_yamlrrr rrr=rs ` @r*rrs>**,33::L77<< .GHL : &Bww||BWW\\,0PQN#$4nEK$$##& <rsY" 88 ,,2/4H?SS*(  U NNMN5N Np/3RB.RB4RBRB% RB  RB  RB,RB RBjU@66665 6 N 6 & 6.66 6F !&#oooo5 o N o & ooooooo ofOSO O "%O 04O AKO  O d zFr,