j,BUddlmZddlZddlZddlZddlZddlZddlmZm Z m Z ddl m Z m Z ddlmZddlmZmZmZmZddlmZddlZddlmcmZddlmcmZddlmcmZddlmcm Z ddl!m"Z"ddl#m$Z$dd l%m&Z&dd l'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.dd l/m0Z0m1Z1m2Z2m3Z3dd l4m5Z5m6Z6dd l7m8Z8m9Z9m:Z:m;Z;mm?Z?ddl@mAZAmBZBmCZCmDZDmEZEmFZFmGZGmHZHmIZImJZJmKZKmLZLmMZMmNZNmOZOmPZPmQZQmRZRmSZSmTZTmUZUmVZVmWZWmXZXmYZYmZZZddl[m\Z\m]Z]m^Z^m_Z_ddl`maZaddlbmcZcmdZdmeZemfZfmgZgmhZhmiZiddljmkZkmlZlerddlmmnZnmoZoedZpGddelZqe dddgZriasdetd<iaudetd<dd&Zv ddd3Zwddd4Zxejydd5Zz dd(dd6dd:Z{dd>Z|dd@Z}ddDZ~ddJZddLZddNZddSZddTZddUZddVZe dWXGdYdZZe dWXGd[d\Ze dWXGd]d^Ze dWXGd_d`Ze dWXGdadbZe2ddcZddfZddiZe dWXGdjdkZddmZddnZddqZddtZddvZdd~ZddZe2ddZddZe3ddZddZddZddZe"jdddZddZe"jdddZddZddZddZddZddZddZddZddZddZddZddZedkr edSdS)) annotationsN) defaultdict namedtuple OrderedDict) dataclassfield)Path)AnyLiteral TYPE_CHECKINGTypeVar) assert_never)cpp) translate)Binding CppSignatureCppSignatureGroupDispatcherSignature NamedCTypeNativeSignatureSpecialArgName)method_with_native_functionnative_function_managerwith_native_function with_native_function_and_indices)gen_aoti_c_shim_files*gen_static_dispatch_backend_call_signature) gen_functionalization_definition"gen_functionalization_registration.gen_functionalization_view_inverse_declaration,gen_functionalization_view_meta_classes_decl,gen_functionalization_view_meta_classes_implGenCompositeViewCopyKernelgen_all_vmap_plumbing)Argument BackendIndexBackendMetadataBaseOperatorNameDEFAULT_KERNEL_NAMESPACEdispatch_device_map DispatchKeyFRAGMENT_NAMESPACESFunctionSchemais_cuda_dispatch_keyis_generic_dispatch_keyis_ufunc_dispatch_keyis_xpu_dispatch_keyLocationNativeFunctionNativeFunctionsGroupNativeFunctionsViewGroup OperatorName OptionalType SchemaKind SelfArgumentSTRUCTURED_DISPATCH_KEYSTensorOptionsArgumentsTypeVariantViewSchemaKind)add_generated_native_functionsgen_composite_functional_kernelgen_composite_out_kernelpre_group_native_functions)SelectiveBuilder) concatMapcontext FileManagermake_file_managermapMaybeNamespaceHelperTarget) YamlDumper YamlLoader)CallableSequenceTc eZdZdfd ZxZS) LineLoaderFcrt||}|jjdz|d<|S)N)deep__line__)superconstruct_mapping start_markline)selfnoderTmapping __class__s V/lsinfo/ai/hellotax_ai/base_platform/venv/lib/python3.11/site-packages/torchgen/gen.pyrXzLineLoader.construct_mappings8''++Dt+<<"o2Q6 )F)__name__ __module__ __qualname__rX __classcell__)r^s@r_rRrRs=r`rR ParsedYamlnative_functionsbackend_indiceszdict[str, ParsedYaml]_GLOBAL_PARSE_NATIVE_YAML_CACHEzdict[str, set[str]]_GLOBAL_PARSE_TAGS_YAML_CACHE dispatch_keyr, device_fmsdict[str, FileManager] default_fmrGreturnc|tfdtjDd|}|S)Nc38K|]\}}||VdSN).0checkdevicerjs r_ z1file_manager_from_dispatch_key..sK  !E65&&       r`)getnextr+items)rjrkrmfms` r_file_manager_from_dispatch_keyr|se      %8%>%@%@          B Ir`Fesobject valid_tagsset[str] ignore_keysset[DispatchKey] | Nonepathstrskip_native_fns_genboolc  t|tstdt|g}t t }|D]}t|t std|t|dtstd|t||d|d td|t fd5tj |||\}} | |tj|| dddn #1swxYwYt|t d} |st!|||D]8\} } t| d d t%| pt'| | | | <9t)|| S) N$Expected 'es' to be a list, but got zExpected to be dict: rVExpected '__line__' to be int: funczMissed 'func' in cddSNzin z: rr)funcslocsr_z*parse_native_yaml_struct..s4344U44r`c>ttjdddiS)NTFrjuse_out_as_primaryexternal device_guardindex)r'r, Undefinedrrr`r_rz*parse_native_yaml_struct..s) $.#   r`TFr) isinstancelistAssertionErrortyperdictrxintr3rFr4 from_yamlappendr' grow_indexerror_check_native_functionsr@rzr/r2re)r~rrrrrsbsermindiceskvrrs @@r_parse_native_yaml_structrsP b$  PNDHHNNOOO!BALTARARB  + +!T"" > !<!.s3333T33r`descrwzTag 'z#' must have a non-empty description) rrrrsetrxrr3rFcopypopadd) r~rrre_inamerrrs @@r_parse_tags_yaml_structrs} b$  PNDHHNNOOO55B   !%% ++S11 H !F1!F!FGG GtQz]++uuU|| 33333 4 4  &&((C775>>D7762&&Drzz$%VT%V%V%VWWW FF4LLL                IsA.D<<E E c|tvr[t|5}tj|t}t ||t|<dddn #1swxYwYt|S)NLoader)r)riopenyamlloadrRr)rfr~s r_parse_tags_yamlrs 000 $ZZ X11Z000B2HRV2W2W2W )$ / X X X X X X X X X X X X X X X ) ..s5AA!A)r loaded_yamltags_yaml_pathr object | Nonec|tvrrt|}|Ct|5}tj|t }dddn #1swxYwYn|}t |||||t|<t|S)Nr)rr)rhrrrrrRr)rrrrrrrr~s r_parse_native_yamlrs 222$^44   d 5qYq444 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5B0H    3 1 1 1 '- +4 00sAAArSequence[NativeFunction]Nonec i}tt}|D];}|||jj<||jjj|<|D]}|j{||j}|,t|jjd|jd|jd|js,t|jjd|jd|jdttj }hd}|jj j D]R}|j|vrGt|jj|jf|vr%td|jd|jjdSd |jvrt|jjd krt|jjd krt|jjjd kr|jjj}|jst|jjd t#|jd|j} t)|| dkrt|jjd|ddS)Nz0 is marked as a structured_delegate pointing to z, but z is missing.zY is not marked as structured. Consider adding 'structured=True' to the delegated operator>z random_.fromfromuniform_r!_has_compatible_shallow_copy_typerzArgument name 'z' in function 'z' is a reserved Python keyword. inplace_viewresize_ resize_as_set_z is marked with tag: inplace_view, but it doesn't follow the naming convention for inplace ops - the codegen expects the base name to have a trailing underscore.Frzw is marked with tag: inplace_view. The codegen expects there to be a corresponding out-of-place view op with the name 'z.' and matching schema, but it didn't find one.)rrrrrstructured_delegaterxr structuredrkeywordkwlist argumentsflat_allrrinplacer)base dunder_methodlen) rfunc_map base_func_mapr delegate_funcPYTHON_RESERVED_KEYWORDSEXCLUSION_LISTarg base_nameout_of_place_base_names r_rrs35HBMdBSBSM 22 !afk&'..q1111 77 ,$LL)>??M$$v{XX,XX454IXXX!+ $v{SS,SS454ISSS$'w~#6#6     6#,  Cx333 $$ch/~EE(o#(oo16;ooo af $ $AFK  I--AFK  L00AFK$%%// (I$ $v{ttt&6y'>&& "=!7899Q>>$v{uu;Duuui77r`scn|dd}|dd}|dd}|dd}|d d }|d d }|d d}|dd}d|dS)z1Convert a python string into a c++ string literal\z\\"z\"z\az\b z\f z\n z\v z\t)replacers r_ cpp_stringr]s $A #uA $A $A $A $A $A $A q888Or`backendslist[BackendIndex]list[DispatchKey]ct|dkrgSd|DtjtjtjtjgzS)Nrcg|] }|j Srr)rj)rsbackends r_ z(static_dispatch_keys..{s===$===r`)rr,CompositeImplicitAutograd%CompositeImplicitAutogradNestedTensorCompositeExplicitAutograd&CompositeExplicitAutogradNonFunctionalrs r_static_dispatch_keysrwsQ 8}} ==H===  1  =  1  > A   r`rr4 backend_indexr'DispatchKey | Nonec|j||r|jS|jr tjS|jr tjS|jr tj S|j r tj SdSrq) r has_kernelrj&has_composite_explicit_autograd_kernelr,r5has_composite_explicit_autograd_non_functional_kernelr&has_composite_implicit_autograd_kernelr4has_composite_implicit_autograd_nested_tensor_kernelr)rrs r_get_static_dispatch_backendrs (M,D,DQ,G,G( )) 1A44 @AAA 1A44 ?A@@ 4r` str | Nonec||jrdSg}|D]G}t||}|3|d|jd|dHd|S)N#include r)manual_kernel_registrationrr root_namelowerjoin)rroutputrrjs r_static_dispatch_ops_headerrs <t F21e<<  # MMVakVVL4F4F4H4HVVV    99V  r` list[str]c4dt|DS)Ncg|]}d|d S)#include rr)rsrjs r_rz1static_dispatch_extra_headers..s3     5,444   r`)rrs r_static_dispatch_extra_headersrs+  0::   r`sig"CppSignature | DispatcherSignaturecpp_sigrcDdd}t|}t|}|D])}|jjtjkr ||}n*t ||}dd|DS)Ninput_bindings list[Binding]rncg}|D]|}|jdkrZtttj|jj|j|j|j}| |g| |}|S)N memory_format)nctyperdefaultargument) rrrr possibly_redundant_memory_formatrrrr r)routput_bindingsbindingspl_mem_format_bindings r_add_spl_memory_format_bindingz5translate_args..add_spl_memory_format_bindings)+% 0 0G|..)0%&G+!#O$-***& &&'=>>>>&&w////r`, c3$K|] }|jV dSrqexprrsas r_rvz!translate_args..s$++QV++++++r`)rrrnr)rrrrrr!rr )rrr% src_bindings goal_bindingsrexprss r_translate_argsr/s $ ((L**,,--M :?nM M M88FFL E N lM 2 2E 99++U+++ + ++r`c :t||}|}t||}||}|r|jr|jnt }|dd}d|d|jd|d|d S)N::nativerwreturn ::();) rrr/ get_kernel cpp_namespacer*rrjr ) rrrrrr.backend_metadata kernel_nsnss r_%generate_static_dispatch_backend_callr;s 9a@@G <<>>D 3 ( (E$//22  & 0 > &&& %   :r * *B QR Q Q=5;;== Q Q Q Q Q Q QQr`c tj|dd}|jr!|jr|j}n|j}|td|}t||}tj dd}|j r+d|dtjd|d|d S|jr+d|dtjd|d|d S|jr+d|dtjd|d|d S|jr+d|dtjd|d|d Sd |d d d |DdS)NFmethodfallback_bindingzExpected cpp_sig to be non-Noner1rwr2r3r4r5z5TORCH_CHECK(false, "Static dispatch does not support z forr&c6g|]}t|jSrr)rrj)rsrs r_rz:generate_static_dispatch_fallback_call..s# A A AC" # # A A Ar`z ");)rfrom_native_functionsymintr has_symintsymint_signature signaturerrr/r*rrr,rr rrrrrrr )rrrgcpp_sigsrrr.r:s r_&generate_static_dispatch_fallback_callrGs !5 %%H z%af''))%+$>??? <<>>D 3 ( (E ! )*b 9 9B/ K```{DJJLL``PT``W\```` @Kmmm{QWWYYmm]ammdimmmm 1K```{DJJLL``PT``W\```` ?Klll{PVVXXll\`llchllllKKK A A A A ABBKKK Kr`c t|dksjrdSfd|D}t|dkrt||dSt|dkrt||Sd|D}d|}jjj}g}g}||d|dkr|d |d |d d |d |dg} |D]H} | d| j d| dt|| d It||} d} d| |d| | d| dS)a  For a given `NativeFunction`, find out the corresponding backend and dispatch to it. If more than one backends exist, fallback to static dispatch by determining dispatch key from inputs. Arguments: sig: A CppSignature or DispatcherSignature for this native function we want to use. f: NativeFunction to generate static dispatch. backend_indices: All available backends. Return: C++ code to call backend-specific functions, e.g., "return at::cpu::add(self, other, scale);" rrwchg|].}|sj|jtv,|/Srq)rrrjr;)rsbrs r_rz#static_dispatch..sP    <<??   ! -"::: ;::r`rUcg|][}t|jts8t|jt6|jjT|j\Srr)rr r:r&ris_tensor_likerr*s r_rz#static_dispatch..$sl aj, / / aj( + +  JO * * , ,  r`r&Nz>DispatchKeySet(c10::computeDispatchKey(dtype, layout, device))z$c10::detail::multi_dispatch_key_set()zDispatchKeySet _dk_set = z | ;z=DispatchKey _dk = c10::highestPriorityBackendTypeId(_dk_set);zcase DispatchKey:::rz  z switch (_dk) { z default: z } ) rr r;rGrr rtensor_optionsrrj) rrrgkeysnative_tensor_args tensor_args tensor_optsstmtssubexprs dispatch_coderfallback connectors ` r_static_dispatchr[sG ?q  A$@ r        D 4yyA~~4S!T!WEEE Ta5c1oNNN)).//K&"1K EH L   bM{MMMNNN LLHuzz(/C/CHHHIII LLPQQQM  KE4FKKKLLL L8aGG L L L    6c1oNNHI ^^E    & &       r`T)frozencTeZdZUded<eeZded<ed d Zd S) RegisterSchemarDselector)default_factoryzdict[str, int] known_tagsrr4rnrc|j|sdSdddt|jDzdz}|dkr%dt t |jdSd}||jvr't|j}||j|<d |d |d }|dt t |jd |j|d S)N{r&c3 K|] }d|V dS)z at::Tag::Nrr)rsrs r_rvz*RegisterSchema.__call__..Ys+KKS0300KKKKKKr`}z{}zm.def(z, {}); rwz const std::vector tags_z = ; z, tags_z); ) r_is_native_function_selectedr sortedrrrrrar)r[rr maybe_tagsidxs r___call__zRegisterSchema.__call__Us}88;; 4TYYKKF16NNKKKKKKcQ 4<<?Js16{{33??? ? t & &do&&C$'DOD !MCMMDMMMJ__Js16{{$;$;__DOTXDY____r`Nrr4rnr) rarbrc__annotations__rrrarrkrrr`r_r^r^Pse!&t! create_zQ_typed_handle() { return c10::Dispatcher::singleton() .findSchemaOrThrow(z::name, z::overload_name) .typed(); } )FTr&dispatchKeySetcg|] }|j Srrrr*s r_rz-ComputeOperators.__call__..s-N-N-Naf-N-N-Nr`cg|] }|j Srrrzr*s r_rz-ComputeOperators.__call__..s5V5V5Vaf5V5V5Vr`r3z static auto op = create_z_typed_handle(); return op.r4r5r)rgr { z } )r from_schemarrunambiguous_namerqrK DECLARATIONr overload_namerrdefn DEFINITIONr rrrrr[r) r[rrrdefnsrudispatcher_exprs_str method_basedispatcher_call method_namefn_bodys r_rkzComputeOperators.__call__ns!-af55v{++-- ;&, , ,0   ((**   /0fk.>   231J  .8AF -D-D   ((E( : :   (( $( ? ?    [F- - - &.2EI )-   E(5  #&)+/99)*-N-Ncmmoo-N-N-NN,,(#/KK+/995V5Vcmmoo5V5V5V+W+W("(K"-!%6666 <!<<<."9"9"9a16"9"9"9r`z c10::SymIntint64_trw inline  { return at::_ops::::call(z); }zZ namespace symint { template >> T)suppress_symint_suffixz ); } } )rrAmanual_cpp_bindingrrC signaturesrr}rrr rBr>functionvariantsdeclrr~) r[r sig_grouprCresultr target_sigr. exprs_str intlike_ts r_rkzComputeFunction.__call__s%: ea.B   V&&(( ''))  C,8@@Jcmmooz/C/C/E/EFFE "9"95"9"9"9::Iz &) % 1:-- &  fk2244>G GP8848((fk2244 >G  r`Nrlrarbrcrrkrrr`r_rrs3 (((! (((r`rc<eZdZUded<ded<ed d Zd S) ComputeTensorMethodrprqrrrrr4rnrc tj|jvrdS|jrt d|j|jjjt d|jtj |d|j }|j tj ur5d}|D]}||dz }|S|j tjurt#|j d}|D]}t%j|j}t)||d}dd |D}|d |jd |d d|jjd|d z }|S)Nz*Method variant cannot be an out function: z#Method variant must have self_arg: Tr=rwz const; )r>r&cg|] }|j Srrr(rs r_rz0ComputeTensorMethod.__call__..rr`rwrzTensor::)prefixz const { return at::_ops::r); } )r>r>rr is_out_fnrrself_argrrArrqrKrrrrrrr}rr rrr~r[rrrrrr.rs r_rkzComputeTensorMethod.__call__s > + +4 6     X !Vaf!V!VWW W 6  $ , !Oqv!O!OPP P%: dQ-A    ;&, , ,F ++-- 3 3SXXZZ2222M ;f/ / /  % % %''))  C,8@@Jcmmooz/C/C/E/EdSSSE "9"95"9"9"9::I  &  ##fk2244>G FF r`Nrlrrrr`r_rrsM::::7777 $$$! $$$r`rc&eZdZeddZdS)ComputeRedispatchFunctionrr4rnrc tj|d|j}d}|D]}t j|j}t||}d dgd|Dz}|d|jd| d d |jj d |d z }|S)NFr=rwr&rxcg|] }|j Srrr(r*s r_rz6ComputeRedispatchFunction.__call__..4s7N7N7N17N7N7Nr`rwrT)rurz ::redispatch(r) rrArrrr}rrrr rrr~rs r_rkz"ComputeRedispatchFunction.__call__(s&: ea.B   ''))  C,8@@Jcmmooz/C/C/E/EFFE #3"47N7N7N7N7N"NOOI  & T**fk2244DM FF r`Nrlrrrr`r_rr&s3 ! r`rcPd|jjjd|jjjdS)Nz{"aten::z", "z"},)rrrrs r_compute_aten_oprEs) Lqv{' L LQV[-F L L LLr`gr5c|jsdSt|j5tj|}tj|}dd|D}|jj}|d}d}|jr |jjnd}|r1g|j |j }d|D}d|D} dd| D} d| d } d |D} d d | D} g}t|D]D\}}d| d|d gz| |dzdz}d|d }| | d}|d|jd|d}d|jd}d| |d|d}g}||dt|D]P\}}||kr|d|jd*|d|jd|jdQ|dd |}|d!|d"|d#|d$Fd |}dd gt| z}d%|d&}d'}d(| d)|d!| d*}nd+}d+}d,|d-|d.|d/|d/|d0|d1 cdddS#1swxYwYdS)2Nr&c3>K|]}|VdSrq)rr*s r_rvz4compute_meta_function_declaration..Qs*44!QVVXX444444r`zat::impl::MetaBasevoidcg|] }|D]}| Srrrr)rs replace_listelems r_rz5compute_meta_function_declaration..]s9$$$% $$@D$$$$r`c@g|]}|jSrr)rupperrsrs r_rz5compute_meta_function_declaration..`s1///&* !!///r`c3"K|] }d|dV dS)zbool z = falseNrr)rsparams r_rvz4compute_meta_function_declaration..cs?88,1''''888888r`z template <>cDg|]}tj||jS))binds)r argument_typerrs r_rz5compute_meta_function_declaration..is:333(TY???333r`rfc3VK|]$}|dd|jV%dS)T strip_ref N)cpp_typerrs r_rvz4compute_meta_function_declaration..nsT33==4=00>>49>>333333r`truerUzprecompute_outrNz return ret;rz z { z z' } z'using meta_return_ty = precompute_out ;meta_return_tyz zG struct TORCH_API precompute_out { z; };rwzstruct TORCH_API structured_z : public r|rPz meta(z); }; )rroutmetarmeta_argumentsr structured_inherits precomputedrvaluesr enumeraterrr)rrargsargs_str parent_class meta_returnrprecomputed_valuesprecomputed_elementsprecomputed_template_parametersprecomputed_template_params_strprecompute_template_decl#precomputed_elements_with_cpp_typesprecomputed_elements_declsetter_methodsirreturn_ty_templates return_ty elem_cpp_tyrE assert_msg assert_stmtconstruction_stmtsjconstruction_blocksetter_methods_declmeta_return_template_paramsmeta_return_typedefprecomputed_decls r_!compute_meta_function_declarationrKs <t  ' 'vvy||(++9944t44444u0  /L +,<Aae''T d ""R;#6#=#=#?#?!Q!Q $$);$$$ //.B/// +/3ii885T888// +(W4S'V'V'V $330333 / ). 33?333)) % N$%9::2 2 4 '+ii3BQB7h5a!egg>?''# E.ADDD A!DMM"N  )OOtyOO;OOO :999 l/Nq/Qll^hlll &(""))Y*=*=*=>>>()=>>GAtAvv*112M2M2M2MNNNN*11C49CCtyCCC#))-888%)YY/A%B%B"%%$,#'))N";"; +/))3>???++ ' #lLg"k"k"k *K )  )  /      #% ! !-9 ! avvvvvvvvvvvvvvvvvvsJ;K&&K*-K*r_rDct|jjj}|ds|drdS|jjjdS||S)N_likenew_F)rrrendswith startswithrrQrg)rr_rs r_needs_backend_selectrsi qv{ D }}W!8!8uv&.u  / / 2 22r`c<eZdZUded<ded<ed d Zd S) ComputeBackendSelectz/Literal[Target.DEFINITION, Target.REGISTRATION]rqrDr_rr4rnrc t||jsdStj|j}t |jd}d|D}tj|j}|}| }d}|j tj ur|r^|jj std|jdd|D} d|d | d } n;|jj rtd |jd |d } d|jd||d| d|jjddd|Dd S|j tjurd|jjd|dSt'|j dS)NTrBcg|]<}t|jt|jj:|=Srr)rr r&rrLr*s r_rz1ComputeBackendSelect.__call__..sT   !*h// 56JO4R4R4T4T    r`z.c10::computeDispatchKey(dtype, layout, device)z'Expected function to have tensor args: r&c3$K|] }|jV dSrqrzr*s r_rvz0ComputeBackendSelect.__call__..s$'K'K1'K'K'K'K'K'Kr`z-DispatchKeySet _dk_set = c10::DispatchKeySet(z() | c10::detail::multi_dispatch_key_set(z); DispatchKeySet _dk_mask = c10::DispatchKeySet(DispatchKeySet::FULL_AFTER, DispatchKey::BackendSelect); DispatchKeySet _dk = c10::impl::computeDispatchKeySet(_dk_set, _dk_mask);z+Expected function to not have tensor args: z)DispatchKeySet _dk = c10::DispatchKeySet(r5z // aten::z C10_ALWAYS_INLINE z { z return at::_ops::z::redispatch( _dk, c3$K|] }|jV dSrqr(r*s r_rvz0ComputeBackendSelect.__call__..s$77af777777r`rzm.impl("aten::z ", TORCH_FN(z));)rr_nativerrrrrr}r.rqrKrhas_tensor_argrr rr~ REGISTRATIONr) r[rr native_sigrSdispatcher_sigrdispatcher_exprsrjrT compute_dks r_rkzComputeBackendSelect.__call__se#At}55 4{16""$QVD999   ))++    -8@@)//11G ;&+ + + " v'6688(J!&JJ#ii'K'K8J'K'K'KKK M.:MMdoMMM 6#2244(NafNNQ PPP &$  FK0022  II77&6777 7 7  [F/ / /JafkJJtJJJ J  % % % % %r`Nrlrrrr`r_rrsO;;;; 8&8&8&! 8&8&8&r`rdatacdt_d d}tjt|t j|dtd S) NcdS)NTrr)r[rs r_rzformat_yaml.. s4r`dumperr rrncP||Srq)represent_dictrz)rrs r_dict_representerz%format_yaml..dict_representer#s$$TZZ\\222r`FgeA)default_flow_styleDumperwidth)rr rr rnr )rLignore_aliasesadd_representerrrdump)rrs r_ format_yamlr sR 7 7J3333{,<=== 9TeJc R R RRr`c|dkrdS|dkrdS t|S#t$r' t|cYS#t$r|cYcYSwxYwwxYw)NrTfalseF)r ValueErrorfloatrs r_pythonify_defaultr1sF{{t gu1vv  88OOO   HHHHH s)! A=A A A AAtr=ct|trt|jSt |dkrdSt j|dddS)NTensorz at::TensorF__placeholder__)mutablerrB)rr8 dynamic_typerrrargumenttype_typer)rs r_rrIsi!\""$AF### 1vv|   5 1%   hjjr`r set[Variant]cdg}tj|vr|dtj|vr|d|S)Nr=r namespace)r>r>rr)r method_ofs r_compute_method_of_yamlrVsTI~!!"""8##%%% r`+tuple[list[dict[str, str]], dict[str, str]]ci}tj|}g}tt|jj|D]\}\}}t |j|tj|d d}|j rG|j |d<|j r$|j ||jj j |j <||||fS)NFr)rrr field_name)r return_namesrziprreturnsrr return_typerrrrrr)rname_to_field_namenamesr!rrrrets r_compute_returns_yamlr'asT*,  Q  EG!#afne"<"<== 9At(00OAe444==??    6 J !C v!! JCD6"16#3#7#:#?@s & &&r`cpp_ar schema_orderkwarg_only_set out_arg_setr#dict[str, str]c.t|jtr&ddd|j|jdd}|j |j|d<|St|jt rtt|jtrt|j||||SdS)Nzat::TensorOptionsFT) annotationr is_nullablerr kwarg_onlyrr)r*r+r#) rr r<rrrr:rr&compute_argument_yaml)r(r)r*r+r#rs r_compute_cpp_argument_yamlr3s%."899 / JJ " "  = $"]C N ENL 1 1   ENH - - $ N%)#1       r`r+r&c b|jrt|jndt|j|j|jt j|ddd}|j 1tt j |j |jd|d<|j|vrd|d<|j|vr#d|d <d|d <|j|vr||j|d <|j }|)|j "t|jd kr |j |d <|S)NrF)rrB)r.rr/rrrrTr0rallocaterrsize)r.rrrr/rrrrrr default_expr is_list_liker6r)r+r)r*r+r#rls r_r2r2s>,-<Ac!,'''T$QV,,v))++!!+z+compute_declaration_yaml..sGGGafGGGr`ch|] }|j Srrrzr*s r_r<z+compute_declaration_yaml..s888a16888r`Fr=c 8g|]}t|dS)Fr1)r3)rsr(r*r#r+s r_rz,compute_declaration_yaml..sI     " )#1       r`c 8g|]}t|dS)Tr1)r2)rsr+r*r#r+s r_rz,compute_declaration_yaml..sI      )#1       r`c ng|]2}tj|dtdddD] }|j 3S)F)r>cpp_no_default_argsfaithfulrBhas_tensor_options)rr rr)rsr+r%s r_rz,compute_declaration_yaml.. sj      #$           r`rz (r&rMc3JK|]}t|jtVdSrq)rr r<r*s r_rvz+compute_declaration_yaml.. s/ M MqJqz#9 : : M M M M M Mr`r operator_namerr category_overriderw schema_stringaten::rschema_order_cpp_signatureschema_order_argumentsr)moder python_moduler!ris_factory_methodabstractr)with_gilF) deprecatedFhas_math_kernel)r'rrflat_kwarg_onlyrrrArErrJr returns_typer!rr anyr>r>rrrrrr rFrrLr is_abstractrr)rr!rcpp_argsrschema_order_jit_argumentsrJcpp_schema_order_types cpp_returnsrIrMr*r#r+s @@@r_compute_declaration_yamlrZs"6q"9"9G HGaf&6&FGGGN8816#3#7888K!6 %%I",,..H         I"&af&C&C&E&E!F!F      ,     ,   "16>%@@@IIKKK$/!W!W499=S3T3T!W!W!W M MH M M MMM - N!* ,  SXaf%% & c!&+"233 4 c!&+";<< = )1+G H#'(':'F##B /qv// 0 ) $ )+E F %'= > 0<< =  AO$;bb Q   (0 1 "3 4  ' Q^ ,  !  H I/   r`c|js |jduoM|jtjkp&|jtjkSrq)rrrkindr9 functionalrrs r_"has_autogenerated_composite_kernelr^BsJ L =A1=   ..U!&++--:CU2Ur`dict[DispatchKey, BackendIndex]c tjj}tjjj}tjj}dd|D}djtfd| Dtj hko9fd| Dtj tj hktj ptd}|d|d|d tj|d S) Nr&c3bK|]*}|V+dSrq) no_defaultrr*s r_rvz4compute_registration_declarations..Os6==1,,..======r`rHcDh|]\}}||Srrrrsrrrs r_r<z4compute_registration_declarations..Ts, F F F41aall1oo FQ F F Fr`cDh|]\}}||Srrrdres r_r<z4compute_registration_declarations..Vs,JJJtq!!,,q//JJJJr`)schemadispatchrrr4z); // r) dispatcherrrrSr!rrr rrzr,rrhas_composite_kernelr^jsondumps)rrgrrSrr comment_datas` r_!compute_registration_declarationsrnHsc ?16 " "D*16>::CCEEL   ' 'Dyy=======H#16## F F F F?0022 F F F56 7 JJJJ4466JJJ5A  q-V1STU1V1VWW $ $L  t  h  dj6N6N   r`"provided_op_registration_allowlistlist[str] | Noneop_selection_yaml_pathc||tdd}|t|}|tj|dd}n*|tj|}ntj}|S)NzhBoth provided_op_registration_allowlist and op_selection_yaml_path can NOT be provided at the same time.TF)rrrD&from_legacy_op_registration_allow_listfrom_yaml_pathget_nop_selector)rorqop_registration_allowlistr_s r_get_custom_build_selectorrwis +6 " . 4   26)5$'(J$K$K! ,#J %     +#23IJJ#466 Or`3Sequence[NativeFunction | NativeFunctionsViewGroup]c4d d}tt}|D]}|j}|j}|t jkr\|j}|||vr-td|d|| ||||<|||vr,t|d|| ||||<tt|| S) Nd1dict[ViewSchemaKind | SchemaKind, NativeFunction]rn/list[NativeFunction | NativeFunctionsViewGroup]c|g}tj|vr|tj}|tjd}|tjd}|t||||| |S)N)view view_copy view_inplace) r?aliasingraliasing_inplacer9r]rr6extendr)rzrr~rrs r_maybe_create_view_groupzEget_grouped_by_view_native_functions..maybe_create_view_groupsBD  "a ' '55011D55!@$GGLj3T::I LL('!-     QXXZZ    r`zDuplicate schema kind z in z already in )rzr{rnr|) rrrview_signatureview_schema_kindr? non_aliasingr\rrRrrEr)rfrgrouped_by_viewsrrg view_kindr\s r_$get_grouped_by_view_native_functionsrsQ. D44&&(($%$6 3 3 36;;==D'///$XTXX7G7O7T7T7V7VXX./ V $T * *,V444$ OO.>v.F.K.K.M.MOO34 V $Y / /  13C3J3J3L3LMM N NNr`/Sequence[NativeFunction | NativeFunctionsGroup]c dd}t|}tt|t|S)Nrz dict[SchemaKind, NativeFunction]rnrc.tj|}|}td|Dr1t dt |t |S|gS)Nc3(K|] }d|jvVdS) generatedN)rrsrs r_rvzJget_grouped_native_functions..flatten_pre_group..s)==Q;!&(======r`zVGenerated NativeFunctions should have been grouped into NativeFunctionsGroup objects: )r5 from_dictrTrrr)rzr%s r_flatten_pre_groupz7get_grouped_native_functions..flatten_pre_groups ! *1 - - 9==!((**===== $H59!((**5E5EHH ## #3Jr`)rzrrnr)rCrrEr)rfrpre_grouped_native_functionss r_get_grouped_native_functionsrs\"$>>N#O#O #T*F*M*M*O*O%P%PQQ  r`)native_function_decl_gengrouped_native_functionsrJCallable[[NativeFunctionsGroup | NativeFunction, BackendIndex], list[str]]dict[str, list[str]]ctt}|D]}t}t}|D]\}}||} | r2| j} |||| nt} t|dkrtd|d|||  ||||S)NrUz6Codegen only supports one namespace per operator, got z from ) rrrrzr6r7rr*rrr) rrgrns_grouped_kernelsrnative_function_namespaces dispatch_keysrj backend_idxr8rs r_get_ns_grouped_kernelsrs30;4/@/@ %%(UU" )8)>)>)@)@   %L+*55a88  5,: !!,///*..y99994 -..22$M5MM=JMM y ) 0 0((K88      r`rc Fg}d}|D]\}}t|dd}ttj|}|d|jd||d|jd ||S)Nrrw namespace_str entity_name max_level ) rzrJrrfromkeysrprologuer epiloguesplit)r declarationsnewlinerkernels ns_helperordered_kernelss r_8get_native_function_declarations_from_ns_grouped_kernelsrs!LG06688   7##   {3G<<==    o       E'NN      r`cFt|||}t|S)a Generate kernel declarations, in `NativeFunction(s).h`. :param grouped_native_functions: a sequence of `NativeFunction` or `NativeFunctionGroup`. :param backend_indices: kernel collections grouped by dispatch key. :param native_function_decl_gen: callable to generate kernel declaration for each `NativeFunction`. :return: a list of string, from the string with all declarations, grouped by namespaces, split by newline. rrgr)r)rr)rrgrrs r_ get_native_function_declarationsrs? 0!9'!9 D-   r`%NativeFunction | NativeFunctionsGrouprc  ||}|r^d|jvrUt|tr |jjn|jjj}td|d|jd|jd|r|jntS)Nr1zThe kernel for function z with dispatch key z has a namespace z% and it's not ending with '::native'.) r6r7rr4rrr]rrjr*)rrr8 func_names r_get_kernel_namespacer(s#--a00 J.>.LLL%a88 TAFKKal>O>T  ey e e!,!9 e e/= e e e   +;X&&@Xr`r{rocmrBskip_dispatcher_op_registrationgen_dispatch_helpersc g} tttttt} d} tj|t j|||d} tj|t j|||d} tj|t j|||d}|D]}t|| dd | | | |t|tr|jn |jj}|| vrtt| <| | ||D]t!dkrt#d| D]?}| |sd|d d | | |d z @| |d fd | | S)NrrrBclass_method_namerrrr1rwr)rz TORCH_LIBRARY_IMPL(r&z , m) { z }zRegisterDispatchDefinitions.inicpjjrdnddS)Nrw) ns_prologue ns_epiloguedispatch_anonymous_definitions"static_init_dispatch_registrationsdeferred_dispatch_registrationsdispatch_namespacedispatch_namespaced_definitions)rrr )anonymous_definitionsrjkernel_namespacens_definitionsrregistration_bodyrsr_rz1get_native_function_definitions..sR#,#5#,#56K(77;+""*79*6*<*<*>*>7EFV7W  r`)rrrdestRegisterDispatchKeyrKNAMESPACED_DEFINITIONANONYMOUS_DEFINITIONrrrrrr4rr]rrJr substitute_with_templater)r{rrjrr_rrBrr definitions registrationsrns_gen anonymous_genreg_genrrrrrrrs ` ` @@@@@r_get_native_function_definitionsr<s K+6t+<+T  M*:; ; ;.9$.?.?M* +&' 299 GAJJ    +   ~./ 0 0A 5 5 #2BCCC &'78  I !129=  """!-"" \\- 01)<=="""      ' '1            eGnn!    & r`c g}tt}d}tj|tj||dd|} |D]]} t | |d|} ||  | | ^| D]\} } t| dkrt| dd } ttj| }| d| jd||d| jd ||S) NrF)rrrrBrrrrwrr)rrrrrKNAMESPACED_DECLARATIONrrr rrzrrJrrrr rr)rrjrr_rrBrrrrrrrrrs r_get_namespaced_declarationrs!L/:4/@/@G  #% (-   D&  (1+FFFNN l((**   9%,, DGG    16688    7 w<<1   ##q   {3G<<==    o       E'NN      r`schema_selectortuple[list[str], str]c tt}|D]"}||j|#d}g}d}|D]k\}}tt t ||} |dkr| }8|}d} |tvrdnd} |d| d|d| | d z }l||fS) NrwatenrTORCH_LIBRARY_FRAGMENT TORCH_LIBRARYrr4z , m) { z };) rrrrrzrIr^r-r ) rfrns_native_functionsnative_functionschema_registrationsaten_schema_registrationscustom_namespacerrschema_registrations_bodytabtorch_library_macros r_(get_native_function_schema_registrationsrs5 .s.,.04:r>rrsfns r_rz*gen_aggregated_headers.. s-2;)F)F)F)F)Fr`cg|]}|v| Srrrr)rsrmethod_native_functionss r_rz*gen_aggregated_headers..s+###3J)J)J)J)J)Jr`MethodOperators.hc vgttttjdS)NrrMethodOperators_includesMethodOperators_declarationsrrIrorKr)rrsr_rz(gen_aggregated_headers..sL(*,0$*8K, --   r`z Operators.hc xdgttttjdS)Nz!#include r)Operators_includesOperators_declarationsr)non_method_native_functionsrsr_rz(gen_aggregated_headers..$sO#F"G&*$*8K0 ''   r` Functions.hcztdgtttdS)Nz#include )rFunctions_includesFunctions_declarations)rrrIrrfrsr_rz(gen_aggregated_headers..3sN-J#..$A"A&*#%%$''   r`)rrgzNativeFunctions.hcdgdS)Nz%#include )NativeFunctions_includesNativeFunctions_declarationsrr)rsr_rz(gen_aggregated_headers..Fs)P(Q,8  r`rFunctions_inl.h>DispatchKeyFunctions.hc(tdSN)rjinline_headersrrj inl_headerssr_rz(gen_aggregated_headers..Ts$' $5$5&1r`Functions_inl.hDispatchKeyFunctions_inl.hc fgtddS)NT)rrjrr_rrB)!DispatchKeyFunctions_inl_includesr dispatch_namespaced_declarations)r r)rgrjrrr_sr_rz(gen_aggregated_headers..\sJ9;*6*<*<*>*>8R1I%1$3L$A!)!# 999  r`)writerr|write_with_template)rfrrrr_rgrrkrrrr{rrjrrrs`````` ` @@@@@r_gen_aggregated_headersrsK  LL    %####%### LL  LL  LL 4!9'L LL    & +L*f M M > ) )JLJJJK  " ",,,(     " "000,           " B=r`ops_fmc tt|D]"} | j| #tt} |D]$} | j| | %D]\|ddfd|ddfd| g}d|Dtdkr|d d fd t||tj |d dfddD]#\| dfd$| D]2|vr gD]\| g}tttj|tj|| ddd|tdkru|dddfdt#||}dd|ddfd|ddfd ~4| d!fd"dS)#Nz_ops.hz Operator.hc tdttttjiS)Nrrr) functionsrsr_rz*gen_per_operator_headers..sF("....s#=.A$$$r`r_ops.h>)static_dispatch_ops_headersoperator_includesfunction_definitions)rrIr)r!rrsr_rz*gen_per_operator_headers..sy/3" 00&I4%H%H%H(,'))!))r`cJg|] }t|t|j|!Srr)rr5rrs r_rz,gen_per_operator_headers..sE   "233 9;     r`rz_meta.hzNativeMetaFunction.hcLdtttiS)Nmeta_function_declarationsr)structured_functionssr_rz*gen_per_operator_headers..s,0$ =?S33r`rz _native.hzNativeFunction.hcrddngdS)Nrz_meta.h>)extra_includesnative_function_declarationsrr)r is_structuredrsr_rz*gen_per_operator_headers..s+rr)rsrsuffixs r_rz>gen_per_operator_headers....s9)))<$;;;;)))r` _declarations)rhrR)categoryfunctions_by_root_namer<sr_rz*gen_per_operator_headers..sa&&&)))) &'='B'B'D'D E E)))***B r`TFrrz _dispatch.hzDispatchKeyFunction.hcdS)N)rrrr)rrsr_rz*gen_per_operator_headers..s*<8Dr`rrrrc(tdSrrrsr_rz*gen_per_operator_headers..s #L 1 1"-r`rrcBfdtDgdS)Nc"g|] }d|dd S)rrr rr)rsrrs r_rz>gen_per_operator_headers....s>666R$QQ1CQQQ666r`)rrr)rh)dispatch_namesrsr_rz*gen_per_operator_headers..sE&86666 &~ 6 666657 r`rc`tdDgdS)Nc3ZK|]&\}}td|Dd|dV'dS)c3<K|]}tj|jvVdSrqrrs r_rvzGgen_per_operator_headers.....(s+IIw~4IIIIIIr`rr%N)rT)rsrr!s r_rvz=gen_per_operator_headers....%sb//#D)IIyIIIII/3d333//////r`r)rhrz)r?sr_rz*gen_per_operator_headers..$sI(.//'='C'C'E'E///)) -/   r`)rrr rrzrrxrrr#compute_native_function_declarationrr rErrKrr|)rfrrr_rgrrkrrrrrgrouped_functions_by_root_namegroupgrouped_functionsr{r>rrjrDrr!r?rr0rr,r<s ` @@@@@@@@@@@@r_gen_per_operator_headersrLms ?J$>O>O88r|,33B7777 D#*;;&t,33E::::17799I I i""OOO         ""KKK          ,;>>tRHH  '    011A5    & &   &   8%6+%)%M   ""             &  OOO       &<< ~ - - )//115;;==  OD) > B B4 L L ,' 55 !#*.8=&    L<  A%%  ! !$ ' ' '  & &99,999'    ,L*f M MF FFF   ( ( ( $          , , , (       B LL         r`core_fmper_operator_headersc |rt|||| | | | |  nt||||| | | |  |dfd|dfd|dfd|d fd dfd }|d|dfd }|d|dS)N) rfrrr_rgrrkrrrr) rfrrrr_rgrrkrrrz TensorBody.hc ttttjttttjdS)N)rqrr)tensor_method_declarationstensor_method_definitions)rrIrrKrrr sr_rzgen_headers..`s*.'%18K% ++*.'%08K% **  r`zRedispatchFunctions.hcZdtttiS)Nfunction_redispatch_definitions)rrIrrfsr_rzgen_headers..xs. -t2446FGG00 r`zRegistrationDeclarations.hc$dfdDiS)Nregistration_declarationsc0g|]}t|Srr)rn)rsrrgs r_rz1gen_headers....s3***2!_EE***r`rr)rgrfsr_rzgen_headers..s3 '****)*** r`zVmapGeneratedPlumbing.hc"tSrqr$rUsr_rzgen_headers..s+@AQ+R+Rr`rnr,c&t}t}D]}|t|jjj||jjjj|d|jD|hdz}ddt|Dddt|DdS)Nc3$K|] }|jV dSrqrz)rsrs r_rvzAgen_headers..gen_aten_interned_strings..s$PPcPPPPPPr`> orandnotxorbitorcomplor_eqand_eqbitandnot_eqxor_eqz \ cg|]}d|d S)z_(aten, rMrrrsrs r_rzBgen_headers..gen_aten_interned_strings..$>>>#D###>>>r`cg|]}d|d S)z_(attr, rMrrrhs r_rzBgen_headers..gen_aten_interned_strings..rir`) aten_symbols attr_symbols) rrrrrrupdaterJr rh)attrsr$rrfs r_gen_aten_interned_stringsz.gen_headers..gen_aten_interned_stringss%%$ Q QD IIc$).-.. / / / IIdin). / / / LLPPTY-M-M-O-OPPP P P P P     $LL>>u >>>$LL>>u >>>    r`zaten_interned_strings.hcLddtiS)Nenum_of_valid_tagsz, )r rh)rsr_ gen_tags_enumz"gen_headers..gen_tags_enums"$uzz&2D2D'E'EGGr`z enum_tag.h)rnr,)rLrr)rfrrrrr_rgrMrrkrrrrrNrorrs`` ` ` r_ gen_headersrs/s$  -%= 3+!')  -%=(C 3+!')  MM     0 LL     LL$      LL!#R#R#R#R" " " " " " H MM+-FGGGHHHHHH MM, .....r` view_groups"Sequence[NativeFunctionsViewGroup]aoti_fm cpu_vec_fmforce_schema_registrationupdate_aoti_c_shim aoti_backendsset[DispatchKey | None]extend_aoti_c_shimc  !"#$%&'()*+,-d}rd}| D]+$t$| | &|r dK"$% fd }ndK$ fd }$"tt}D]E}t|tr|jn |jj}|||Ft$ %$tj k(t$r|nddtj"||(rtj"ngd}dL"$&(fd }&d $ddd|$tjkrdnd|dhD]ފ''jjrt)$s'jjjj*$tjurS&| urt/d&d*dd$'fd| d*dd'*fd$tjur)d#rd#&d*dd #$'*fd!t/d"$d#&-t5|||||$dMfd% }| d&|}|rt9j}t=|'\!+| d(!+fd)dNd-}| d.|fd/d0tAid1hd23| d4tB|d5tB|d6fd7dOd9)dP)fd: }ttz} d;tEd<D,d=tEd>D-| #,-fd?D| d@| ||dhdAB| dCfdD| dEfdF| dG)fdH| dIfdJdS)QNz#include #include #include #include z#include #include #include #include rnrc6g}D]}d}|rd}nht|tr0tfd|Drd}n#|jrt jt jfvrd}|s| d|j dt jkr| d|j dvr!| d|j ddtt|S) NFTc3BK|]}|VdSrqrd)rsrrs r_rvz=gen_source_files..operator_headers..sJEE9; 0044EEEEEEr`r _native.h>r;rr ) rrr5rTr!rr,Metarrr rhr)headersr is_registeredrrjrrrs r_operator_headersz*gen_source_files..operator_headerss1""A$)M$//22-(, $A';<< -EEEE?@{{}}EEEBB -)- -,#(#J;++)- (! NN#P#P#P#PQQQ$&MNN 'MQ['M'M'MNNN#~55`!+``@R```c'll+++r`cdg}tjkr|dvr|dd|S)Nz!#include z#include rr)r,rr)rrjrs r_rz*gen_source_files..operator_headers sZ>?;#UUUNN#@AAA>11NN#Q\#Q#Q#QRRRr`rw)extra_cuda_headersexternal_backend_headersdispatch_headers ops_headersdispatch_helpersgnfrrc :dt|gd iS)Ndispatch_definitionsT) r{rrjrr_rrBrr)r)rrrjr{rrr_rs r_"register_dispatch_key_env_callablez.register_dispatch_key_env_callable6 sA'(G.1U!- -%4S)= ) ) )  r`Registerz.cppzRegisterDispatchKey.cppc|jSrqr )xs r_rz"gen_source_files..K sQ[r`rrUr)key_fn env_callable num_shardsbase_env sharded_keysz,Expected fm to be cpu_fm for DispatchKey.CPU UfuncCPU_z UfuncCPU.cppcttjtjdS)N)meta_declarationnative_declarationnative_definitions)rrrHcompute_ufunc_cpu)rgrjrsr_rz"gen_source_files..\ sD,Ma,P,P.2.V|<///3.DQ.G.G r`UfuncCPUKernel_zUfuncCPUKernel.cppc2tjdS)N)rr)rcompute_ufunc_cpu_kernel)rrsr_rz"gen_source_files..g s $.2.KA.N.Nr`z%#include z$#include UfuncCUDA_z.cuz UfuncCUDA.cucttjtjdS)N)r cuda_headersrrr)rrrHcompute_ufunc_cuda)rgrrjrrsr_rz"gen_source_files..s sJ $(4,Ma,P,P.2.V|<///3.Ea.H.Hr`z unrecognized z for ufunc)rvrzrfrgrrryr|c  fdD}d|Dttttj|ttttj|dS)Nc4g|]}t||Srr)r)rsrr_s r_rz@gen_source_files..gen_backend_select.. s9   -A"h-O-O    r`c$g|] }d|jdS)rr%rrs r_rz@gen_source_files..gen_backend_select.. s3@B;bl;;;r`)r!backend_select_method_definitions%backend_select_function_registrations)rrIrrKrr) relevant_fnsrfr_s r_gen_backend_selectz,gen_source_files..gen_backend_select s    )   FR26():HEE|22 6:(). s-.*+*.%&BB%   r`r@NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGrouprc|jSrqr)rs r_key_funcz"gen_source_files..key_func s |r`z Operators.cppchd|jdgttj|gdS)Nrr;r)rr)r rorKrr$s r_rz"gen_source_files.. sV!Hr|!H!H!H I %4G! ! r`r>rrr)rrrrrz Functions.cppzTensorMethods.cppzATenOpList.cppcLdtttiS)Naten_ops)rrIrrUsr_rz"gen_source_files.. s" Xo7GHHII r`rcBt|trLd|jjdd|jjdg}|j#|d|jjdd|jjdgz }|St|t rd|jjdd|jjdd|jjdd|jjdg}|j#|d|jjdd|jjdgz }|j #|d|j jdd|j jdgz }|Sd|jdd|jdgS)Nrrr%) rr6r~r rr5r]rrr)rrs r_gen_op_headersz(gen_source_files..gen_op_headers s a1 2 2" Caf&6BBB?af&6???G{&K!+*?KKKH!+*?HHHN / 0 0 Hal&<HHHEal&<EEEAaeoAAA>aeo>>> G y$I!)*=IIIF!)*=FFFy$I!)*=IIIF!)*=FFFN>ak===:ak::: r`c||t|t|tjdS)N)rfunc_definitionsfunc_registrations)rrr,r)rrgrr_s r_functionalization_env_callablez8gen_source_files..functionalization_env_callable sU*>!,, @!!#E EF##   r`c(i|]}|jj|Srrrrrs r_ z$gen_source_files..% s.:::   Q:::r`cDt|Srqrr!rs r_rz"gen_source_files..' sT!++--%8%8r`c(i|]}|jj|Srrrrs r_rz$gen_source_files..) s+444 Q444r`cDt|Srqrrs r_rz"gen_source_files..* sD4G4Gr`c3TK|]"}|jjv|jjv|V#dSrqr)rsrstructured_mapview_maps r_rvz#gen_source_files.., sL 6;n , ,H1L1L 1L1L1L1Lr`zRegisterFunctionalization.cpp>rrrfunc_add_back_views_definitions!func_add_back_views_registrations)rrrrzFunctionalInverses.hcHdttfdiS)Nview_inverse_declarationsc$t|Srq)r rr_s r_rz4gen_source_files....F sL !r`)rrIr_rtsr_rz"gen_source_files..C s? ' **  r`zViewMetaClasses.hcHdttfdiS)Nview_meta_declarationsc$t|Srq)r!rs r_rz4gen_source_files....T J8UVWWr`rrErsr_rz"gen_source_files..Q s5 $dWWWW'' r`zViewMetaClasses.cppcttfdttdS)Nc$t|Srq)r"rs r_rz4gen_source_files....` rr`)view_meta_implementations op_headersr)rr_rtsr_rz"gen_source_files..] sO)-WWWW** yEEFF  r`zCompositeViewCopyKernels.cppc 0dDdDzttttjttt ttt dS)Ncg|]<}dd|j|jgn |j|jgD=S)rc3<K|]}d|jd|jdVdS)r_ops.h> #include ....| sb    B!+BB+,+ BBB      r`)r rr~rsrs r_rz6gen_source_files....{ sn     %&K$7afak=R        r`cpg|]3}dd|j|j|jfD4S)rc3RK|]"}|d|jv d|jd|jdV#dS)Nrrrr)rr rs r_rvz@gen_source_files..... sl}AF)B)B B!+BB*++BBB*C)B)B)B r`)r rrr]rs r_rz6gen_source_files.... sc     iALA    r`)r#CompositeViewCopyKernel_Definitions(GeneratedCompositeFunctional_Definitions!GeneratedCompositeOut_Definitions)rrIr#r,rrArB)rgrrtsr_rz"gen_source_files..z s  %     5   .48.''N   4 49=3/99 26,/22Q. . r`)rnr)rrrnr)rnr)rrrnr)rrrnr)rrrnr)$r|rrrr4rr]rrr r,rr/rgen_registration_headersgen_registration_helperswrite_sharded_with_templateCPUrufunc_inner_loopr1rrrrCUDArrrDrur write_shardedrrrEr).rfrrrtr_rrgrvrMrwrrkrrrrxrNrryrzr|rrns_grouped_native_functionsgrouped_native_functionrregister_dispatch_key_base_envrrrrr all_groupsrrrrjrr{rrrrrrrs.``````` `` ` @@@@@@@@@@@@@r_gen_source_filesrs0&  $$ &cc +L*f M M 2 & ,& ,& ,& ,& ,& ,& ,& ,& ,& ,& ,T       ( 5 &1$&7&7#'? S S #5~FFB'11,7A  ( 2 9 9:Q R R R R ..4466 KM M  $L11#"4"4(* $ =3T!!,+--(-m<<<* * &"            " && )| ) ) ) % $((;(KO;;qq301 ' -, O, OA5) 1F|1T1T <$).D{..V##()WXXX&&****"   ..0d000(!111F J#IL&&****"    %%M\%M%M%MNNN B#)'$?---           ( LL,.@AAAO >*;== 1)?   ! LL                ,-J#..    '4 LL$''' MM%t,,, MM    %%%%N        $ ())D-- ::88:UVV:::N44 )*G*G U U444H!  '3       LL     LL         LL         : LL&. . . . . . 11111r`c:|dfddS)NzDeclarations.yamlc6tdDS)Nc,g|]}t|Srr)rZrs r_rz;gen_declarations_yaml.... s!SSSQ5a88SSSr`)r rUsr_rz'gen_declarations_yaml.. s SSBRSSSTTr`)r)rrfs `r_gen_declarations_yamlr s4 LLTTTTr`r cXttjS)z If you're depending on torchgen out-of-tree, you can use the root to figure out the path to native_functions.yaml )r __file__parentresolverrr`r_get_torchgen_rootr s >> ( ( * **r`cN#$%tjd}|dddd|dd d |d d d|dd d|ddddd|dddd|dd d|dd d|dd d |d!d d"|d#d$d%d&'|d(d)d* |d+d,d%d-'|d.d/d%d0'|d1d2d d3|d4d5d d6|d7td%gd8gd8d9:|d;d d<|d=d d>|%t %j%j}tj %j d?}tj %j d@}dAdBl m }tjtjtjtjtjtjtjtjh}tjtjdh}t/}tjtjtjh#%js%jr5|#|tjn(|##fdC|D|dd<%js%jr?|tj |tj nM|tj tj |vr ||!tj =%j"sM|tjtj|vr ||!tj=%j#r%fdD|D}tI|||}tJ|} |j&|j'c} $tQ| } dE| D} tS| } dF| D}%j*dG}tW|,dHdHI%j*dJ}tW|,dHdHI%j-}tW|,dHdHIt]%|K}t]%L}t]%L}t]%L}t]%|K}t]%|K}dM|i}%jrt]%L|dN<g}%j/rJ$fdO%j/D}%j/D]/}tj0|}||vr||0dP%j1vrgtedoidQ| dR| dS| dT|dU|dV|dW$dX|dY|dZ|d[|d\|d]|d^|d_%j3d`%j4da%j5db%j6dc%jdd|de%j7df%j1vr(tq| | | | ||$||||||%j3%j5gdh%j1vrts| |i%j:rtW%j:;}|j<}|j=}|djf|dkf|dlf|dmfgdn|>DzD]<\}} | |z}!|j?| |zz }"|@|!t|";dSdS)pNzGenerate ATen source files) descriptionz-sz --source-pathz!path to source directory for ATenz aten/src/ATen)helprz-oz--output-dependenciesz:output a list of dependencies into the given file and exit)rz --dry-run store_truez5run without writing any files (still updates outputs))actionrz--per-operator-headersz2generate separate headers per operator in ATen/opsz-dz --install-dirz --install_dirzoutput directoryzbuild/aten/src/ATenz--aoti-install-dirz--aoti_install_dirz%output directory for AOTInductor shimz(torch/csrc/inductor/aoti_torch/generatedz--rocmz=reinterpret CUDA as ROCm/HIP and adjust filepaths accordinglyz--mpsz'Generate MPS registration code when setz--xpuz'Generate XPU registration code when setz--mtiaz(Generate MTIA registration code when setz--op-registration-whitelistz--op_registration_whitelist*zfilter op registrations by the whitelist (if set); each item is `namespace`::`operator name` without overload name; e.g.: aten::empty aten::conv2d ...)nargsrz--op-selection-yaml-pathz--op_selection_yaml_pathaDProvide a path to the operator selection (for custom build) YAML that contains the information about the set of selected operators and their categories (training, ...). Each operator is either a full operator name with overload or just a bare operator name. The operator names also contain the namespace prefix (e.g. aten::)z--backend-whitelistz--backend_whitelistzRfilter dispatch backend by the whitelist (if set), e.g.: CPU CUDA QuantizedCPU ...z--static-dispatch-backendz--static_dispatch_backendz?generate static dispatch code for the specific backend (if set)z!--skip-dispatcher-op-registrationz!--skip_dispatcher_op_registrationz0Avoid registering operators into the dispatcher.z--force-schema-registrationz--force_schema_registrationz}force it to generate schema-only registrations for all ops, includingthose that are not listed on --op-registration-whitelistz --generate)rsourcesdeclarations_yamlzGenerate only a subset of files)rrchoicesrrz--update-aoti-c-shimzUpdate AOTInductor C shim after adding an entry to inductor_fallback_ops in torchgen/aoti/fallback_ops.py. WARNING: Do not use this unless you are sure what you are doing!!!z--extend-aoti-c-shimaThis Flag indicates the generation of c shims for out-of-tree ATen ops,which is an extension to the In-tree ATen op c shims. This flag needs to be combined with---source-path=--aoti-install-dir=/extend default is torch/csrc/inductor/aoti_torch/generated/extendWARNING: Do not use this unless you are sure what you are doing!!!znative/native_functions.yamlznative/tags.yamlr)rcg|]}|v| Srrrr)rsrMPS_KEYSs r_rzmain..g s#JJJ!8I8IA8I8I8Ir`c^g|])}t|st|jv'|*Srr)r0rbackend_whitelist)rsroptionss r_rzmain..y sI   &q)) .1VVw7P-P-P -P-P-Pr`c<g|]}t|t|Srr)rr5rs r_rzmain.. s8###z!=Q/R/R# ###r`c<g|]}t|t|Srr)rr6rs r_rzmain.. s9 a1 2 2 r`z/coreT)parentsexist_okz/ops)r install_dir)rcudaxpucDg|]}tj|Srr)r,parse)rskeyrgs r_rzmain.. s9    K-c22 3   r`rrfrrrtr_rrgrvrMrwrrkrrrrxrNrryrzr|r)rfrrrrr_rgrMrrkrrrrrNr)rfrrwcpu_vec_core_ops_c"g|] \}}||df S)rrr)rsru device_fms r_rzmain.. s( S S S+<69iF & S S Sr`rr)AargparseArgumentParser add_argumentr parse_argsrwop_registration_whitelistrqosrr  source_pathtorchgen.modelrr,rrrrrrrMTIArMPS SparseMPS SparseCsrMPSmpsryrmrr XPUrmtiarrrirfrgrrrr mkdiraoti_install_dirrHstatic_dispatch_backendr generaterrrxrNrr|rsroutput_dependenciesrrstemrzr write_outputs)&parserr_native_yaml_pathrrrrzr parsed_yamlrrfrr!native_functions_with_view_groupsrtcore_install_dirops_install_dirr"rMrrwcuda_fmrrvrkrr dp_key depfile_path depfile_name depfile_stemr{rvarnamerrrgrs& @@@r_mainr4 s~  $1M N N NF   0     I   D    A    %   4:    L   6   6   7 %%-  ""M *  ## N   ++ ?   %%C   ;;;;;; .  M  M   !!G()&H w||G$79WXXW\\'"57IJJN,,,,,, -9-: N  M%%K!6 8PQH{Kg0Kh'''+/****8$$$JJJJ}JJJ aaa{Dg0D;?++++/**** ((( ?m + +m11+/BBC <E ()))  } , ,m11+2BCCD      "   $$4nkRRK.~>J$#&o <>> 5.0&+    6   2 + +C &s++F^++""6***G$$$   -- %=%= )D(C $  X  !4 3  ,O G G "z 6 "z (- *>  '.&G&G! "")!=!=# $-4,S,S% & '99' ((-) * '99+  0G$$$-!%=(C 3+!')!(!=    $g.../?OOOO" 1G788@@BB #( #( RL  $ g  V   T S @P@P@R@R S S S T 1 1JB |+G&&<*?@D   Wc$ii 0 0 0 0 1 1  1 1r`__main__)rjr,rkrlrmrGrnrG)Nr}F) r~rrrrrrrrrrnre)r})r~rrrrnr)rrrnrrq) rrrrrrrrrrrnre)rrrnr)rrrnr)rrrnr)rr4rr'rnr)rr4rrrnr)rrrnr)rrrrrnr)rrrr4rr'rnr)rrrr4rgrrnrr)rr5rnr)rr4r_rDrnr)rrrnr)rrrnr)rr=rnr)rrrnr)rr4rnr) r(rr)rr*rr+rr#r,rnr) r+r&r)rr*rr+rr#r,rnr)rr4rnr)rr4rnr)rr4rgr_rnr)rorprqrrnrD)rfrrnrx)rfrrnr)rrrgr_rrrnr)rrrnr)rrrgr_rrrnr)rrrr'rnr)r{rGrrrjr,rr'r_rDrrrBrrrrrrnr)rrrjr,rr'r_rDrrrBrrnr)rfrrrDrnr)rfrrrrrrrr_rDrgr_rrGrkrlrrrrrrrnr)rfrrrrrr_rDrgr_rrGrkrlrrGrrrrrrrnr) rfrrrrrrrrrr_rDrgr_rMrGrrGrkrlrrGrrrrrrrNrrnr),rfrrrrrrtrur_rDrrrgr_rvrGrMrGrwrGrrGrkrlrrrrrrrxrrNrrrryrrzr{r|rrnr)rrGrfrrnr)rnr )rnr) __future__rr functoolsrkrr collectionsrrr dataclassesrrpathlibr typingr r r r typing_extensionsrrtorchgen.api.dispatcherapiritorchgen.api.metartorchgen.api.nativertorchgen.api.structuredr torchgen.destr torchgen.apirtorchgen.api.translatertorchgen.api.typesrrrrrrrtorchgen.contextrrrrtorchgen.gen_aoti_c_shimrr#torchgen.gen_functionalization_typerrr r!r"r#torchgen.gen_vmap_plumbingr%rr&r'r(r)r*r+r,r-r.r/r0r1r2r3r4r5r6r7r8r9r:r;r<r=r>r?#torchgen.native_function_generationr@rArBrC!torchgen.selective_build.selectorrDtorchgen.utilsrErFrGrHrIrJrKtorchgen.yaml_utilsrLrMcollections.abcrNrOrPrRrerhrmrir|rrcacherrrrrrrrr/r;rGr[r^rorrrrrrrr rrrr'r3r2rZr^rnrwrrrHrrrrrrrrrLrsrrrr4rarrr`r_rPs5 """""""  <<<<<<<<<<((((((((777777777777****** ,,,,,,,,, $$$$$$$$$,,,,,,,,,,,,,,, =<<<<<8 ?>>>>>76666666322222222 GCLLBZ '9;L&M N N :<;;;;577777,,0 % 1#1#1#1#1#h( ////,01 !&!% 111111@====@    4     (      , , , ,FRRRR$KKKK0O0O0O0Of@ 0 >@ 0 4(____H,,,,`!=!=!=!=HvvvvrDG/G/G/G/Tppppf++++z1z1z1z1z  zDFFFFFr`