tjdZddlmZddlZddlZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl Z ddl Z ddlZddlZddlZddlmZddlmZddlmZmZmZmZddlZddlmZddlZddlmZddl m!Z!dd l"m#Z#dd lm$Z$dd l%m&Z&dd l'm(Z(m)Z)d dl*m+Z+d dlm,Z,m-Z-m.Z.erddl/m0Z0m1Z1ddl2m3Z3ddl4m5Z5e j6e7Z8edZ9edZ:e:;ZgZ?dZ@dZAejBjCjDEFddZAdGde?DZ@gdZHGddZIdkd ZJd!ZKGd"d#eLZMGd$d%ZNejOdkd&ZPd'd(dld+ZQd'd(dld,ZRdkd-ZSdmd0ZTGd1d2eUZVdnd6ZW dodpd=ZX dqd'd'd>drdBZYdsdDZZdtdIZ[dudJZ\ dqd'd'd>dvdMZ]dwdSZ^dxdVZ_e_ej`Zae_ejbdWZce_dZde_d'Zee_d'ZfGdXdYZgGdZd[ZhGd\d]Zi dydzdgZjd{djZkdS)|a Debug utilities for TorchDynamo compilation and execution. This module provides various debugging tools and utilities for TorchDynamo, including: - Minification support for reducing test cases while preserving bugs - Input/output handling via InputReader and InputWriter for reproducible testing - Accuracy checking between original and compiled models - Neural network module string conversion via NNModuleToString - Profiling tools and system information collection - Buck build system integration for Meta-internal testing Key classes: - InputReader/InputWriter: Handle serialization of model inputs/outputs - NNModuleToString: Converts nn.Modules to string representations - BuckTargetWriter: Manages Buck build system integration ) annotationsN)Counter) import_module)Anycast TYPE_CHECKINGTypeVar)Tensor) rand_strided)normalize_path_separator)is_float_dtype)StorageWeakRef)ContentStoreReaderContentStoreWriter)config) clone_inputs get_debug_dir warn_once)CallableSequence)tqdm)UntypedStorageTztorch._inductor.config)z1//caffe2/torch/fb/sparsenn:sparsenn_operators_gpuz-//caffe2/torch/fb/sparsenn:sparsenn_operatorsz///deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpuz+//deeplearning/fbgemm/fbgemm_gpu:sparse_opszfbcode:// cg|]}d|d S)ztorch.ops.load_library("").0xs W/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/torch/_dynamo/debug_utils.py r%Ts$TTTA?!???TTT)buck2runz@mode/dev-nosanc(eZdZd dZddZddd Zd S)BuckTargetWriterfilenamestrreturnNonec(tjtj|\|_|_|jdd|_|jddd|j|_|j|jdd|_|jdd|_|j}||dddd}d|d |j|_ dS) Nz.pyr/.zfbcode.zfbcode/r:) ospathsplitabspathsubdirpy_filereplacetargetfind cmd_line_path)selfr+tmps r$__init__zBuckTargetWriter.__init__[s$&GMM"'//(2K2K$L$L! T\l**5"55 {**344DDt{DD IdinnY7799: IabbM k#((9%%''(,5#55 55r&c ddtD}tjd|jd|jdt d|d|jd S) Nrcg|]}d|d S)z "z",r r!s r$r%z*BuckTargetWriter.build..js$#J#J#J!$5$5$5$5#J#J#Jr&za load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") python_binary( name="z", srcs = ["z"], compile = False, deps = [ "//caffe2:torch", "//caffe2:libtorch", "//caffe2/functorch:functorch", "//triton:triton", "z", ], cpp_deps = [ z ], main_module = "z", par_style = "xar", ) )join extra_depstextwrapdedentr;r9 cur_targetr5)r>extra_cpp_depss r$buildzBuckTargetWriter.buildis#J#Jz#J#J#JKK  ;    l        "I#      r&T print_msgbool list[str]cntj|jd}t |d5}||dddn #1swxYwYt|jgz}|r.t dd||S)NTARGETSwzFFound an example that reproduces the error. Run this cmd to repro - %s ) r4r5rCr8openwriterIBUCK_CMD_PREFIXr=logwarning)r>rJ target_filefd cmd_splits r$rRzBuckTargetWriter.writesgll4; :: +s # # #r HHTZZ\\ " " " # # # # # # # # # # # # # # #$t'9&::   KKX##   s(A**A.1A.N)r+r,r-r.r-r,)T)rJrKr-rL)__name__ __module__ __qualname__r@rIrRr r&r$r*r*ZsU 6 6 6 6    4       r&r*r-r,ctjtd}|)t jdt j}tj|stj |d|S)Nminifierz /minifier_T)exist_ok) r4r5rCrtempfile gettempdirgetpassgetuserexistsmakedirs)r5s r$ minifier_dirrfst 7<<  4 4D |%''FF7?3D3DFF 7>>$  ) D4(((( Kr&ceZdZdS)UnsupportedNNModuleErrorNrZr[r\r r&r$ririDr&riceZdZddhZhdZddhZdZejj ejj ejj ejj ejj ejjejjejjejjejjejjejjejjejjejjejjejjejjejjejjejjgZed d Z ed!dZ!ed"dZ"ed#dZ#ed#dZ$ed#dZ%ed#dZ&ed$dZ'eddd%dZ(eddd&dZ)dS)'NNModuleToStringz0torch.ao.quantization.fake_quantize.FakeQuantizezAtorch.ao.quantization.fake_quantize.FusedMovingAvgObsFakeQuantize> #torch.ao.nn.qat.modules.conv.Conv1d#torch.ao.nn.qat.modules.conv.Conv2d#torch.ao.nn.qat.modules.conv.Conv3d5torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBn1d5torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBn2d5torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBn3d7torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvReLU1d7torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvReLU2d7torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvReLU3d9torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBnReLU1d9torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBnReLU2d9torch.ao.nn.intrinsic.qat.modules.conv_fused.ConvBnReLU3dz%torch.ao.nn.qat.modules.linear.Linearz8torch.ao.nn.intrinsic.qat.modules.linear_relu.LinearReLU)fbgemmx86qnnpackonednngmtorch.fx.GraphModuler-rKct}|D]4\}}t|||5t |dkrt d|dSdS)Nrz2Cannot safely convert some modules to strings - %sFT)setnamed_childrenrm_module_constructor_stringaddlenrTrU)r~ cant_convert_modules r$can_convert_to_stringz&NNModuleToString.can_convert_to_stringsuu **,, ) )IAv::6BBJ  ((( |  q KKDl   5tr&rtorch.nn.Moduler,cXt|jdt|jSNr1)typer[rZ)rs r$_module_type_namez"NNModuleToString._module_type_names(v,,)CCDLL,ACCCr&type_ type[object]c$|jd|jSr)r[rZ)rs r$ _type_namezNNModuleToString._type_names"55U^555r& str | Nonect|}|tjvrdS|j}t t |j}t t |j}t|ddr |dz}|dzdz}dt t|d|d|d|j d |j d |j g}t|dr|d |jt|d r|d |j|dd|dS)N reduce_rangeFrz observer=z quant_min=z quant_max=dtype=zqscheme=z is_dynamic=z reduce_range=ch_axiszch_axis=(, ))rmrfake_quant_modulesactivation_post_processrint quant_min quant_maxgetattrrrdtypeqscheme is_dynamichasattrappendrrrC)r module_typeobserverrrargss r$_fake_quant_constructorz(NNModuleToString._fake_quant_constructorsi&88@@ .A A A41f.// f.// 8^U 3 3 * NI!A )I E(33DNNCC E E & & & & & & %V\ % % )v~ ) ) /&+ / /   8^ , , C KKA(=AA B B B 8Y ' ' 9 KK78#377 8 8 822 $2222r&ct|dd}tjD]`} tjj|}n#t$rY4wxYwt|t|krd|dcSadS)Nqconfigz6qconfig=torch.ao.quantization.get_default_qat_qconfig(r) rrmdefault_qat_qconfig_backendstorchao quantizationget_default_qat_qconfigAssertionErrorrepr)rrbackenddefault_qconfigs r$_qat_qconfig_constructorz)NNModuleToString._qat_qconfig_constructors&)T22'D  G "'("7"O"OPW"X"X!    G}}_ 5 555%%%%6 ts$A AAc Jt|}|tjvrdSt|}|dSt |jt |jd|jd|jd|j d|j d|j d|j dud|j g }t|dd}|0|d |jd |jd |jg|||d d |dS)Nz kernel_size=zstride=zpadding=z dilation=zgroups=bias=z padding_mode=bnzeps=z momentum=z freeze_bn=rrr)rmrqat_conv_modulesrr in_channels out_channels kernel_sizestridepaddingdilationgroupsbias padding_moderextendepsmomentum freeze_bnrrC)rrrrrs r$_qat_conv_constructorz&NNModuleToString._qat_conv_constructors]&88@@ .? ? ?4";;FCC ?4 # $ $ $ % % 16- 1 1 'fm ' ' )v~ ) ) + + + 'fm ' ' -FKt+ - - 3F/ 3 3  VT4 ( ( > KK%26%%/ //5!155    G22 $2222r&c0t|}|tjvrdSt|}|dSt |jt |jd|jdu|g}|dd|dS)Nrrrr) rmrqat_linear_modulesrr in_features out_featuresrrC)rrrrs r$_qat_linear_constructorz(NNModuleToString._qat_linear_constructor6s&88@@ .A A A4";;FCC ?4 # $ $ $ % % -FKt+ - -    22 $2222r& module_strct|tjvrdS td|ddn#t$rYdSwxYwdS)NFzmod = z execT)rrm safe_reprscompile SyntaxError)rrs r$_can_emit_module_constructorz-NNModuleToString._can_emit_module_constructorGsf <</: : :5  )Z))?F C C C C   55 ts4 AAFallow_unsafe_reprrc t|}t||r|St|p3t|pt|}||S|r|SdSN)rrmrrrr)rrrconstructor_strs r$rz+NNModuleToString._module_constructor_stringRs&\\  8 8 L L    4 4V < < @55f== @77??   &" "   tr&cddlm}d}tjd}|D]\}}t ||}|td|t| d}||n!t| d} | | j j dkr |d| j d }||d zd |d |d z }|j D]\} } | | t krFddlm} | jt kr t)d| jdt dt+| } nUt-j| r!dt1| jd| jd} n dt1| jd| jd} | j j dkr | d| j d } ||d zd| d| dz }|jD]X\}}|d}|j j dkr d|j d}dt1|jd|j|d} ||d zd |d | d z }Y|||jdd z }|S)Nr) _addindent z from torch.nn import * class Repro(torch.nn.Module): def __init__(self) -> None: super().__init__() rz!Cannot convert module to string: cpuz.to("rrzself.z = r) PRINT_OPTSzPRINT_OPTS.threshold (z() must be >= MAX_CONSTANT_NUMEL_INLINE (rz torch.randn(z, dtype=ztorch.randint(1, size=zself.register_buffer('z', z) rz , device=""ztorch.nn.Parameter(torch.randn(z))rg)torch.nn.modules.modulerrErFrrmrrinext parametersbuffersdevicer_buffersitemsnumelMAX_CONSTANT_NUMEL_INLINEtorch._tensor_strr thresholdrrris_floating_pointlistshaper _parameterscode)r~rrtab model_str module_namerr example_paramexample_tensor buffer_namebufferr tensor_str param_nameparam maybe_devices r$convertzNNModuleToString.convertdss666666O   $#4#6#6 I I K)DD*;EJ!.BBB !!2!2!4!4d;;M!, &..**D11  )n.C.HE.Q.Q *JJ1FJJJ C!GHH+HH*HHH HII#%;#4#4#6#6   K~||~~!:::888888'*CCC(S1ESS6OSSS"&\\ (00 WD,>,>WW WWW YT&,-?-?XXXXX}!U** *BBBBB 7QQ+QQ*QQQ II"$!5!5!7!7 H H J}L| E));EL;;; s4 ;L;LssV[VascosssJ C!GGG*GGGGG GII  27A..2222 r&N)r~rr-rK)rrr-r,)rrr-r,)rrr-r)rrrr,r-rK)rrrrKr-r)r~rrrKr-r,)*rZr[r\rrrrrnnLinearConv1dConv2dConv3d BatchNorm1d BatchNorm2d BatchNorm3d LayerNormDropoutSoftmaxReLUGELUIdentity MaxPool2d EmbeddingTanhConvTranspose1dGLULSTMFlattenAdaptiveAvgPool2dr staticmethodrrrrrrrrrrr r&r$rmrms[:K    0B$J                            "+J0   \ DDD\D666\6333\32   \ 333\3>333\3 \>C\"GLLLLLL\LLLr&rmcjtjsdSd} tjjet jddg}|d}d d|D}||dz }n|dz }n!#tt j f$r|d z }YnwxYwtd ttjD}|d z }|D]\}}|d |d |dz }|dz }|S)Nz:# torch.cuda.is_available()==False, no GPU info collected z# CUDA Info: nvccz --versionrrc&g|]}|dkd|dS)r#  r )r"ss r$r%z-_cuda_system_info_comment..s%RRRq!r''{A{{{'''r&z'# Not searching for nvcc on ROCM setup z# nvcc not found c3TK|]#}tj|V$dSr)rcudaget_device_name)r"is r$ z,_cuda_system_info_comment..sC*+ ""1%%r&z# GPU Hardware Info: rz : r)rr is_availableversionhip subprocess check_outputdecoder6rCFileNotFoundErrorCalledProcessErrorrrange device_countr)rcuda_version_outcuda_version_linescomment gpu_namesnamecounts r$_cuda_system_info_commentr.sx : " " $ $MLL!I * =  $)6 7LMM !1!8!8!:!:!@!@!F!F ggRR4FRRRSSG G 'II C CI z< =***)) */4UZ5L5L5N5N/O/OI**I ((.. e-$--5----  I sA;B B>=B>F) stable_outputr/rKc|rdSgdgdddgdfd fd tjD}|d Dd |}t d|dS)zl Generate a string configuration for environment variables related to Dynamo, Inductor, and Triton. z+# env var omitted due to stable_output=True)TORCHDYNAMOINDUCTORTRITON)TRITON_LIBDEVICE_PATHTRITON_PTXAS_PATHTRITON_LIBCUDA_PATHTORCHDYNAMO_REPRO_AFTERTORCHDYNAMO_REPRO_LEVELkeyr,r-rKcLtfdDovovS)Nc3 K|]}|vV dSrr )r"stringr:s r$rz;generate_env_vars_string..filter..s'77&# 777777r&)any)r: allow_listrepro_env_vars skip_lists`r$filterz(generate_env_vars_string..filtersC 7777J777 7 7 *9$ *>) r&c fg|]-\}}|d|d|ddd.S)z os.environ['z'] = ''r)r:)r"r:valuerBs r$r%z,generate_env_vars_string..s\ C 6#;;@@@EMM#s$;$;@@@r&c3"K|] }d|dV dS)zos.environ.pop('z', None)Nr )r"r:s r$rz+generate_env_vars_string..s.SSS83888SSSSSSr&rz import os z )r:r,r-rK)r4environrrrCr )r/ config_lines config_stringr?rBr@rAs @@@@r$generate_env_vars_stringrJs=<<:::JUUUI/1JKN        ***,,L SSNSSSSSSIIl++M #%%%%   r&c &ddl}ddl}|rdS|jjj}d|jjd|jjd|j jd|d S)Nrz*# config omitted due to stable_output=Truez~import torch._dynamo.config import torch._inductor.config import torch._functorch.config import torch.fx.experimental._config r) torch._functorch.configtorch._inductor.configfx experimental_configcodegen_config_dynamor _inductor _functorch)r/rexperimental_configs r$generate_config_stringrVs""""!!!!<;;(/7FFHH  $$&&   &&((  ''))      r&cZtjtdS)Nzminifier_launcher.py)r4r5rCrfr r&r$get_minifier_repro_pathrXs 7<< (> ? ??r&contentsr.ct}td|tr!t | t |d5}||ddddS#1swxYwYdS#t$r2}tdtd||d}~wwxYw)NzWriting minified repro to: %srOrzCould not write to ) rXrTrUuse_buckr*rRrQOSError exceptionNotImplementedError)rYminified_repro_pathrWes r$helper_for_dump_minifyras#133KK02EFFF6,--33555V %s + + r HHX                      VVV b!"M8K"M"MNNTUUVs<B#B9 BB  B B B C-C  CceZdZdS) AccuracyErrorNrjr r&r$rcrcrkr&rcexample_inputs Sequence[Any] list[Any]ct|}tt|D]H}t||tjr&||||jI|S)z This clone inputs is different from utils clone_input. In case of minifier, all the tensors are leaf tensors while creating a new graph. So, we set the requires_grad field w/o checking the leafness of the tensor. )rr&r isinstancerr requires_grad_ requires_grad)rd cloned_inputsidxs r$clone_inputs_retaining_gradnessrmst !00MS(())QQ mC(%, 7 7 Q #  - -nS.A.O P P P r&r~rronly_fwd disable_clonercdddlm}m}m}t j|}|st |}t|dr|dt|ddr ||n||}|r|S||r||}| |||d|S)z Runs a forward and possibly backward iteration for a given mod and args. When disable_clone is True, we will use args as-is without cloning. This is higher fidelity but we may destroy the args in the process. r)collect_resultsreduce_to_scalar_lossrequires_bwd_pass zero_gradT _boxed_callFN) testingrqrrrscopydeepcopyrmrrtrbackward) r~rrnrorqrrrsoutlosss r$run_fwd_maybe_bwdr|(sSRRRRRRRRR r  B 5.t44r; Tb-77 F""T(((RRYC $$S))  ?2sD$ / //r& require_fp64 ignore_non_fpopt_gmr~rcddlm}t|||}d}tjr t t j|t|\} } t| | |}n;#t$r.|rtdt dYnwxYw t|||} n+#t$rt dYdSwxYw||| |tjd|} | S) aa Check two models have same accuracy. require_fp64: if True, raise an error if we unable to calculate the fp64 reference ignore_non_fp: if True, do not compare outputs which are not floating point. This is mostly useful for the minifier (which wants to avoid quantizing floating point error into integer/boolean error) r)sameNzfCould not generate fp64 outputs, workaround with torch._dynamo.config.same_two_models_use_fp64 = FalsezCould not generate fp64 outputszWhile minifying the program in accuracy minification mode, ran into a runtime exception which is likely an unrelated issue. Skipping this graph.T)tol equal_nanr)utilsrr|rsame_two_models_use_fp64 cast_to_fp64rwrxrm Exception RuntimeErrorrTrUr]repro_tolerance) r~rrdrnr~rrreffp64_ref fp64_model fp64_examplesrespassings r$same_two_modelsrHs]" B 9 9CH & ; ;(4 b!!#B>#R#R)) %J )]HMMHH ; ; ; "| KK9 : : : : :  ; AA   $   ttd    "# G Ns$AA++5B#"B#'B99$C! C!modelc|jjD]}|jdkr|jtjjjjurt|j dkr$tdt|j t|j dr9|j dtj kr|j dtj f|_ |jdkrU|jd}|9t|r*t!|j}tj |d<||_!|j||S)N call_functionrz"Expected node to have 2 args, got rrr)graphnodesopr;ropsprimsconvert_element_typedefaultrrrr float64kwargsgetdictlint recompile)rnoder new_kwargss r$cast_dtype_args_to_fp64rs/ !)) G & & uyCKKK49~~""$ITYIIdil++ : !  0M0M!Yq\5=9 7o % %KOOG,,E ^E%:%: !$+.. &+m 7#(  K OO Lr&r torch.dtypeinputs&tuple[torch.fx.GraphModule, list[Any]]cddlm}|}tjkrt |}|fd|}||fS)Nr)tree_mapct|tjr)|r|n|Sr)rhrr rto)r#rs r$zcast_to..sA a & &+,+>+>+@+@!$$u+++ r&)torch.utils._pytreerrrrr)rrrrs` r$cast_torsx-,,,,, HHUOOE  (.. X      F &=r&c8ttj||Sr)rrr)rrs r$rrs 5=% 0 00r& compiler_fnACallable[[torch.fx.GraphModule, list[Any]], torch.fx.GraphModule]c |tj|t|}t|||||| S#t$rt dYdSwxYw)Nr}zWhile minifying the program in accuracy minification mode, ran into a runtime exception which is likely an unrelated issue. Skipping this graphF)rwrxrmrrrTr])r~rdrrnr~r compiled_gms r$backend_accuracy_failsrs!k M"  >~NN  #    %'         #   uusAA$A+*A+r%torch._prims_common.StrideType | Nonertorch._prims_common.ShapeTypetorch._prims_common.StrideTypec2||ntj|Sr)rmake_contiguous_strides_for)rrs r$_stride_or_defaultrs '66U-Nu-U-UUr&dCallable[[T | None], T]cfdS)Nc||nSrr )r#rs r$rz_mk_defaulter..s!-QQQr&r )rs`r$ _mk_defaulterrs . . . ..r&rcVeZdZddZdddd dZd!dZd"dZd#dZd#dZd$dZ d%dZ dS)&NopInputReaderr-r.cd|_dS)Nrtotal)r>s r$r@zNopInputReader.__init__s  r&Nr dtype_hint storage_hashrnbytesrr)torch._prims_common.DeviceLikeType | Nonertorch.dtype | Nonec&|xjdz c_dSNrr)r>rrrrs r$storagezNopInputReader.storages a r&rrrtorch.Tensor | NonecdSrr r>rrs r$tensorzNopInputReader.tensor r& int | NonecdSrr rs r$symintzNopInputReader.symintrr&r,r,cdSrr r>r,s r$constzNopInputReader.constrr&cdSrr rs r$ unsupportedzNopInputReader.unsupportedrr& device_type device_indexcdSrr )r>rrs r$ generatorzNopInputReader.generator rr&script_class_namec&|xjdz c_dSrrr>rs r$opaquezNopInputReader.opaque s a r&r-r.) rrrrrrrrr-r.)rrrrr-r)rrrrr-rr,r,r-r.)rr,rrr-r.rr,r-r. rZr[r\r@rrrrrrrr r&r$rrs=A)-                     r&rcveZdZ d1ddd2d Zddd d3dZ d1dddddd4d#Zdd$d5d'Zd6d*Zd6d+Zd7d.Z d8d0Z dS)9 InputReaderN)pbarsave_dirrr tqdm | Noner-r.c|td|t|nd|_g|_||_dS)Nz0no save_dir specified, will generate random data)rTrUrstorerr)r>rrs r$r@zInputReader.__init__sI   KKJ K K K5=5I'111t !  r&rrrrrrrrrc|j|jdt|}t|}|j[|Y |j|}||jkr!td||j|S#t$rYnwxYwtd|d||j zf}td|}t||||S)Nrzdevice mismatch: %s != %szcould not load z , generating random data insteadr)rupdate_device_or_default_dtype_or_defaultr read_storagerrTrUr$ritemsizerr untyped_storage)r>rrrrrrrs r$rzInputReader.storage!s 9 I  Q   #F++&z22 : !l&> *11,??W^++KK ;VW^TTT%     RLRRRSSS:..0#D666E6:v>>NNPPPs B BB)storage_offsetrrjis_leafrrrrrrrrrj bool | Nonermetadatar torch.Tensorc t||}t|}t|}t|}t |}t jg||j|} t j5| ||||dddn #1swxYwY|st j 5| t j } dddn #1swxYwYt j5| ||||dddn #1swxYwYt j j| |kr7t!dt j j| d|t j| ||j| | S)Nr)rrrj) memory_formatz/Tensor leaf status mismatch: safe_is_leaf(t) = z , expected )r_storage_offset_or_defaultr_is_leaf_or_default_requires_grad_or_defaultrrrno_gradset_ enable_gradclonepreserve_format _subclasses meta_utils safe_is_leafr_utilsset_tensor_metadatarr) r>rrrrrrjrrts r$rzInputReader.tensor>s$F%8883NCC!%((%g..1-@@ L eGN-   ]__ ; ; FF7NE6 : : : ; ; ; ; ; ; ; ; ; ; ; ; ; ; ; ?"$$ A AGG%*?G@@ A A A A A A A A A A A A A A A ? ?wv>>> ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?   ' 4 4Q 7 77 B B V$/<B##B'*B'!C00C47C4D33D7:D7)exprvalr c|j||6t|dsi|_||jt |jdz <|S)N symint_exprsr)rrrrr)r>r r s r$rzInputReader.symintcsW   400 746!48D c$)nnq0 1 r&r,r,c:|jddSrrrrs r$rzInputReader.constk r&c:|jddSrrrs r$rzInputReader.unsupportednrr&rrctj||}tj|}|j||S)N)r)rr Generatorrr)r>rrrgens r$rzInputReader.generatorqs@k<88oV,,,  r&rc:|jddSrrrs r$rzInputReader.opaquewrr&r)rrrrr-r.) rrrrrrrrr-r)rrrrrrrrrrrjrrrrrr-r)r rr rr-rr)rr,rrr-rrrr r&r$rrs%) BF      $=A)- QQQQQQB9= # &*$(%)#######J6: r&rcdeZdZddd$d Zd%d Zd d d d&dZd'dZd(dZd)dZd*dZ d+d!Z d,d#Z d S)- InputWriterF stable_hashrrrrKr-r.cg|_tj|_||_|t ||nd|_i|_dS)Nr)_lines itertoolsr-storage_counterrrr seen_storages)r>rrs r$r@zInputWriter.__init__sW!# (00  # x[ A A A A 9;r&rLc~dg}|d|jD|d|S)Nzdef load_args(reader):c3 K|] }d|V dS)rNr )r"ls r$rz$InputWriter.lines..s(11111111r&zload_args._version = 0)rrr)r>rs r$lineszInputWriter.linessK $  11T[111111 )***r&N) device_hintrrrr%rrrr,c Jt|}|j|}||Sdt|j}d}t dt |krd|}d}|j}|jdkr|td|}td|krd|}| } d} |j *|jjdkr|j |} |j |d| d| ||d ||j|<|S) Nbufrz , dtype_hint=metaz:device_hint must be provided when storage device is 'meta'z , device=z = reader.storage(rr)rrrrrrrrrrrr write_storagerr) r>rr%rwsvmaybe_dtype_hintrrrrs r$rzInputWriter.storagesoO , ,   " "2 & & =H .$t+,, . . T " "&7 &C&C C C=z==   ' ;& "$P!F d # #v - -1v11L '')) : !o&<&AV&K&K:33ODDL  a aL a af a aN^ a a a   "#2r&r,r rc ddlm}m}|||j|j}g}||td|j| sA| tt| td|jkr| d|j|td|ks*| d|t j|}|r1|d|Dt+d|jkr| d|jt jj|}t5d|kr| d ||j d d |tt|jg|zd |zdS) Nr)statically_known_truesym_eq)rr%rrzstorage_offset=c3*K|]\}}|d|VdS)=Nr )r"kr+s r$rz%InputWriter.tensor..s0IIA1 q IIIIIIr&zrequires_grad=zis_leaf=zreader.tensor(r) # )%torch.fx.experimental.symbolic_shapesr.r/rrrrrrrrr,tuplerrrrrget_tensor_metadatarrrrjrrrrrrC) r>r,r r.r/rrtensor_metadatars r$rzInputWriter.tensorslWWWWWWWW,,    AG  $$ F%d!':::AHHJJ G G   0 KKE!((**--.. / / / T " "ag - - KK,,, - - -$$ &t , ,0@0@0B0B B   B KK@!*:*:*<*<@@ A A A,::1==  J KKII1F1F1H1HIII I I I $T * *ao = = KK<<< = = =#.;;A>> t $ $ / / KK.7.. / / /  ii#eAGnn"5"5==>> ?dnn      r&argrc<|jd|dt|t|tt fr|jdt |D]\}}|d|d}t|tjr| ||>t|ttj fr| ||v| |||jddSdS)Nzreader.unsupported(z#) # unsupported type for dumping: z"""[])rrrrhrr5 enumeraterr rrSymIntrr)r>r,r8raname_is r$rzInputWriter.unsupporteds(  X$ X XTRUYY X X   cD%= ) ) & K  u % % %!# 0 01 1a..0KK****C#6770KK****$$VQ//// K  u % % % % % & &r&cH|jd|d|ddS)Nz reader.const(r3z!, filtered out during compilationrrrs r$rzInputWriter.consts<  PD P P P P P     r&r ct|tjrJt|jj}|jj}|jd|d|d|dS|jd|d|dS)Nzreader.symint(z, expr=r3) rhrr=r,rr hintrr)r>r,r expr_strrCs r$rzInputWriter.symints c5< ( ( D38=))H8=D K  VVVxVVPTVV W W W W W K  BBBDBB C C C C Cr&torch._C.Generatorcn|j}|jd|jd|jd|dS)Nzreader.generator(rr3)rrrrindex)r>r,r8rs r$rzInputWriter.generator sL  L L L L Ld L L     r&rcF|jd|d|dS)Nzreader.opaque(r3rA)r>r,rs r$rzInputWriter.opaques0 L,=LLdLLMMMMMr&)rrrrKr-r.)r-rL)rrr%rrrr-r,)r,r,r rr-r.)r,r,r8rr-r.r)r,r,r rr-r.)r,r,r8rEr-r.)r,r,rr,r-r.) rZr[r\r@r$rrrrrrrr r&r$rrsDI ; ; ; ; ; ;&BF)- ######J    @&&&&(    DDDD    NNNNNNr&rrfunc&Callable[[list[Tensor]], list[Tensor]]r sym_shapesdict[str, int] | Nonedefault_sym_shaperdict[str, Any]c ddlm}d|D}d|}t j|}d|d}d|d} d } Gd d } i} |pidfd dfd } |j}|jD]\}}|dkr tj | |}|rP| \}}t| d}||}| ||| |<tj | |}|r!| d| |<dt j|jvr| }|| d<tj||D]`}| \}}}}t| d}||}t#||| ||a| S)a Takes in a function which has been printed with print_readable() and constructs kwargs to run it. Handles Tensor inputs, Symints, and a graph module which might have tensor constants. Consider a function `forward` defined as follows: def forward(self, primals_1: "f32[1001, 6]", primals_2: "f32[s0]", primals_3: "Sym(s0)",): _tensor_constant0: "i64[4190]" = self._tensor_constant0 # Further implementation kwargs = aot_graph_input_parser(forward) forward(**kwargs) r) dtype_abbrsci|]\}}|| Sr r )r"r:rEs r$ z*aot_graph_input_parser..+s+)))!sEs)))r&|z(_tensor_constant\d+): \"(z0)\[\s*(.*?)\s*\]\" = self\.(_tensor_constant\d+)rz)\[\s*(.*?)\s*\]z Sym\((s\d+)\)ceZdZdZdS)/aot_graph_input_parser..TensorContainerz#Container for tensors as attributesN)rZr[r\__doc__r r&r$TensorContainerrU8s----r&rWrr,r-rcntjvpdufdS)NcdS)Nz; not in symbolic_shapes and default sym shape not passed inr )rsr$rz=aot_graph_input_parser..get_sym_int..CsvZZZr&)r_checkr)rrMsym_shapes_dicts`r$ get_sym_intz+aot_graph_input_parser..get_sym_int@sP o % F):$)F Z Z Z Z   ""6+<===r&rrrrr cg}g}t|D]w\}}|}d|vr6 |}||||S|r"|t|x|jr t jn t j}||| }|D]"} t j || #|S)Nr)rr) r<striprrrrrandnzerosrR mark_dynamic) rrresolved_shape dynamic_dimsrdimr constructorrzrrr\s r$ gen_tensorz*aot_graph_input_parser..gen_tensorGs && 4 4FAs))++CczzK$$%%a(((##A&&&&4"))#c((333%*%<Mekk%+ k.fEEE / /A M & &sA . . . . r&,rr>)rr,r-r)rrrrr-r )torch.utils._dtype_abbrsrPrrCvaluesinspect getsource__annotations__researchrr5r6group signaturerfinditersetattr)rIrrKrMrP dtype_map dtype_patternsourcetensor_assignment_regex tensor_regexsym_shape_regexrWrrfrr annotationmatch data_type shape_strrr container attr_namerr\r[s ` ` @@r$aot_graph_input_parserrs*544444))%0%6%6%8%8)))I+"4"4"6"677M t $ $F|M{{{7 777L&O........ F&0&6BO>>>>>>>*&K.[.0088z H    , 33  5#(<<>> Iy)//#..//Ei(E&Jue44F5M /:66  8'K A77F5M "4((333#O%% "v[!8&AA D DE16 .Iy)Q)//#..//Ei(E Iy**UE*B*B C C C C Mr&r+Callable[[T], T]ctjtjtjdfd }dfd }t j||S) z Decorator to cProfile a given function and save the result to disk on process exit. Args: filename: filename to save profile to fnrr-cLtjdfd }|S)Nrrrr-c |i|S#wxYwr)enabledisable)rrrprofs r$wrapperz3profile_to_file..decorator..wrappersH KKMMM r4*6**  s 3A )rrrrr-r) functoolswraps)rrrs` r$ decoratorz"profile_to_file..decoratorsC            r&r.c tjt jddddS)Nz! Wrote profile to z+, view with: snakeviz z ) dump_statssysstderrrRrErF)r+rsr$save_itz profile_to_file..save_itsn !!!  O"*'   r&)rrr-rr)cProfileProfiler4r5r7 expanduseratexitregister)r+rrrs` @r$profile_to_filer{s    Dwrw11(;;<rs{$#"""""   ######444444444444 ##########......@@@@@@......;;;;;;MMMMMMMM9999999999-22222222,,,,,,g!! GCLL - 899  $ $ & &     VJ $.==??GG SWXXJIITTTTTUUM65544444444n     ~   PPPPPPPPf 67<      @5:(@@@@ V V V V     I         00000H 7 777777t0(1111  RVVVV////"M%-00"]<5<#6#677*]1--)M%00#mE**FeeeeeeeehJNJNJNJNJNJNJNJN^(,$( dddddN######r&