o Zhv @s\dZddlZddlZddlZddlZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl Z ddl Z ddlZddlmZddlmZddlmZmZmZmZddlZddlmZddlZddlmZddlmZddlm Z dd l!m"Z"dd l#m$Z$m%Z%d d l&m'Z'd d lm(Z(m)Z)e*e+Z,edZ-edZ.e./Z0e0rddl1Z2gZ3dZ4e0rgdZ3e2j5j6j789ddZ:d;dde3DZ4gdZdZ?GdddZ@eAddd ZBd!d"d#d$ZCd!d"d%d&ZDd'd(ZEd)d*ZFGd+d,d,eGZHd-d.ZIdWd/d0ZJ !dXd!d!d1d2d3ZKd4d5ZLd6d7ZMd8d9ZN !dXd!d!d1d:d;ZOdd?d@d=fdAdBZPdCe-d@eee-ge-ffdDdEZQeQejRZSeQeTdFZUeQdZVeQd!ZWeQd!ZXGdGdHdHZYGdIdJdJZZGdKdLdLZ[ M  dYdNee\ege\efdOe]dPee^e]e_fdQee_d@e^e]eff dRdSZ`dTe]d@ee-ge-ffdUdVZadS)Za 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 N)Counter) import_module)AnyCallableOptionalTypeVar)Tensor) rand_strided)is_float_dtype)StorageWeakRef)ContentStoreReaderContentStoreWriter)config) clone_inputs get_debug_dirTztorch._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:// cCg|]}d|dqS)ztorch.ops.load_library("z").0xrrH/var/www/auris/lib/python3.10/site-packages/torch/_dynamo/debug_utils.py Kr)Zbuck2runz@mode/dev-nosanc@s&eZdZddZddZd ddZdS) BuckTargetWritercCstjtj|\|_|_|jdd|_|jddd|j|_|j|jdd|_|jdd|_|j}||dddd}d|d |j|_ dS) Nz.pyr/.zfbcode.zfbcode/r:) ospathsplitabspathsubdirpy_filereplacetargetfind cmd_line_path)selffilenametmprrr__init__RszBuckTargetWriter.__init__c CsDdddtD}td|jd|jdtd|d|jd S) NrcSr)z "z",rrrrrrarz*BuckTargetWriter.build..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+r) cur_targetr%)r.Zextra_cpp_depsrrrbuild`s zBuckTargetWriter.buildTcCsntj|jd}t|d}||Wdn1swYt|jg}|r5t dd||S)NZTARGETSwzFFound an example that reproduces the error. Run this cmd to repro - %s ) r$r%r2r(openwriter7BUCK_CMD_PREFIXr-logwarning)r.Z print_msgZ target_filefdZ cmd_splitrrrr;zs  zBuckTargetWriter.writeN)T)__name__ __module__ __qualname__r1r7r;rrrrrQsrcCsLtjtd}|durtdt}tj|s$tj |dd|S)NZminifierz /minifier_T)exist_ok) r$r%r2rtempfile gettempdirgetpassgetuserexistsmakedirs)r%rrr minifier_dirs  rJc@seZdZejjejjejjejjejj ejj ejj ejj ejj ejjejjejjejjejjejjejjejjejjejjejjejjgZeddZeddZdS)NNModuleToStringcCsLt}|D]\}}t|tjvr||qt|dkr$td|dS)Nrz-We have not tested reprs of some modules - %sT) setnamed_childrentyperL safe_reprsaddlenr=r>)gmZ cant_convert_modulerrrcan_convert_to_strings   z&NNModuleToString.can_convert_to_stringcCsddlm}d}td}|D]+\}}|}t|d}|dur-|jr-|d}||dd|d|d 7}q|j D]X\}} | durKqB| t krcdd l m} | jt ks^Jt| } n t| rvd t| jd | jd } n dt| jd | jd } | jr| d} ||dd|d| d7}qB|j D].\} } | durqd}| jrd}dt| jd | j|d} ||dd| d| d 7}q|||jdd 7}|S)Nr) _addindent z from torch.nn import * class Repro(torch.nn.Module): def __init__(self) -> None: super().__init__() z.cuda()zself.z = r) PRINT_OPTSz torch.randn(z, dtype=)ztorch.randint(1, size=zself.register_buffer('z', z) rz, device="cuda"ztorch.nn.Parameter(torch.randn(z))rK)Ztorch.nn.modules.modulerWr4r5rN__repr__next parametersZis_cuda_buffersitemsZnumelMAX_CONSTANT_NUMEL_INLINEZtorch._tensor_strrZ thresholdreprtorchis_floating_pointlistshapedtype _parameterscode)rSrWtab model_str module_namerUZ module_strZ example_paramZ buffer_namebufferrZZ tensor_str param_nameparam maybe_devicerrrconvertsJ           zNNModuleToString.convertN)r@rArBrdnnZLinearZConv1dZConv2dZConv3dZ BatchNorm1dZ BatchNorm2dZ BatchNorm3dZ LayerNormZDropoutZSoftmaxZReLUZGELUZIdentityZ MaxPool2dZ EmbeddingZTanhZConvTranspose1dZGLUZLSTMZFlattenZAdaptiveAvgPool2drP staticmethodrVrrrrrrrLs6  rLc CstjsdSd}z!tddg}|d}ddd|D}||d7}Wnttj fy:|d 7}Ynwt d d t tj D}|d 7}| D]\}}|d |d|d7}qQ|d7}|S)Nz:# torch.cuda.is_available()==False, no GPU info collected z# CUDA Info: Znvccz --versionrrcSs g|] }|dvrd|dqS))r#  r)rsrrrrs z-_cuda_system_info_comment..z# nvcc not found css|] }tj|VqdSN)rdcudaZget_device_name)rirrr s  z,_cuda_system_info_comment..z# GPU Hardware Info: ruz : rv)rdryZ is_available subprocess check_outputdecoder&r2FileNotFoundErrorCalledProcessErrorrrangeZ device_countr`)rlZcuda_version_outZcuda_version_linescommentZ gpu_namesnamecountrrr_cuda_system_info_comments&   rF) stable_outputcsT|rdSgdgdfddfddtjD}d|}d |d S) zl Generate a string configuration for environment variables related to Dynamo, Inductor, and Triton. z+# env var omitted due to stable_output=True)ZTORCHZDYNAMOZINDUCTORZTRITON)ZTRITON_LIBDEVICE_PATHZTRITON_PTXAS_PATHZTRITON_LIBCUDA_PATHcstfddDovS)Nc3s|]}|vVqdSrxr)rstringkeyrrr{sz;generate_env_vars_string..filter..)anyr) allow_list skip_listrrfiltersz(generate_env_vars_string..filtercs*g|]\}}|rd|d|dqS)z os.environ['z'] = ''rrrvalue)rrrr s z,generate_env_vars_string..rz import os z )r$environr`r2)rZ config_linesZ config_stringr)rrrrgenerate_env_vars_strings  rc Cs\ddl}ddl}|r dS|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) Ztorch._functorch.configZtorch._inductor.configZfxZ experimental_configZcodegen_config_dynamorZ _inductorZ _functorch)rrdZexperimental_configrrrgenerate_config_string,s   rcCstjtdS)Nzminifier_launcher.py)r$r%r2rJrrrrget_minifier_repro_path@src Cst}td|trt|zt|d}||WdWdS1s)wYWdStyF}z tdt d|d}~ww)NzWriting minified repro to: %sr8rz(Could not write to {minified_repro_path}) rr=r>use_buckrr;r:OSError exceptionNotImplementedError)contentsZminified_repro_pathr?errrhelper_for_dump_minifyDs    &  rc@s eZdZdS) AccuracyErrorN)r@rArBrrrrrSsrcCsBt|}tt|D]}t||tjr||||jq |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. )rrrR isinstancerdrZrequires_grad_ requires_grad)example_inputsZ cloned_inputsidxrrrclone_inputs_retaining_gradnessWs rc Csddlm}m}m}t|}|st|}t|dr|dt|dr(||n||}|r0|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_gradTZ _boxed_callN) testingrrrcopydeepcopyrhasattrrZbackward) rSargsonly_fwdZ disable_clonerrroutZlossrrrrun_fwd_maybe_bwdds   r require_fp64 ignore_non_fpc Csddlm}t|||}d}tjr:ztt|t|\} } t| | |}Wnt y9|r2t dt dYnwzt|||} Wnt yQt dYdSw||| |tjd|d} | 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)ZtolZ equal_nanr)utilsrrrZsame_two_models_use_fp64 cast_to_fp64rrr Exception RuntimeErrorr=r>rZrepro_tolerance) rSZopt_gmrrrrrrefZfp64_refZ fp64_modelZ fp64_examplesresZpassingrrrsame_two_modelssB     rcCs|jjD]Q}|jdkr5|jtjjjjkr5t |j dksJt |j dr5|j dtj kr5|j dtj f|_ |jdkrU|j d}|durUt |rUt|j }tj |d<||_ q|j||S)NZ call_functionrYrrrh)graphnodesopr+rdopsZprimsZconvert_element_typedefaultrRrr float64kwargsgetdictZlintZ recompile)modelnoderhZ new_kwargsrrrcast_dtype_args_to_fp64s       rcsBddlm}|}tjkrt|}|fdd|}||fS)Nr)tree_mapcs"t|tjr|r|S|Srx)rrdrretorrhrrs  zcast_to..)Ztorch.utils._pytreerrrdrr)rhrinputsrrrrcast_tos    rcCsttj||Srx)rrdr)rrrrrrsrcCsLz|t|t|}t||||||d WSty%tdYdSw)NrzWhile minifying the program in accuracy minification mode, ran into a runtime exception which is likely an unrelated issue. Skipping this graphF)rrrrrr=r)rSrZ compiler_fnrrrZ compiled_gmrrrbackend_accuracy_failss$  rstrideztorch._prims_common.StrideTypergztorch._prims_common.ShapeTypereturncCs|dur|St|Srx)rZmake_contiguous_strides_for)rrgrrr_stride_or_default srdcs fddS)Ncs|dur|SSrxrrrrrrsz_mk_defaulter..rrrrr _mk_defaulters rcpuc@s6eZdZd ddZdddddZdd Zd d ZdS) NopInputReaderrNcCs d|_dS)Nrtotal)r.rrrr1!s zNopInputReader.__init__device dtype_hintcCs|jd7_dS)Nrr)r. storage_hashnbytesrrrrrstorage$szNopInputReader.storagecOdSrxrr.rrrrrtensor'zNopInputReader.tensorcOrrxrrrrrsymint*rzNopInputReader.symintrNr@rArBr1rrrrrrrr s   rc@sLeZdZd ddddZdddddZ d dddddd d Zd d ZdS) InputReaderN)pbarcCs8|dur td|durt|nd|_g|_||_dS)Nz0no save_dir specified, will generate random data)r=r>r storerr)r.save_dirrrrrr11s   zInputReader.__init__rcCs|jdur |jdt|}t|}|jdur=|dur=z|j|}Wn ty-Ynw||jkr;t d||j|St d|||j f}t d|d}t |||| S)Nrzdevice mismatch: %s != %sz1could not load %s, generating random data insteadrg)rupdate_device_or_default_dtype_or_defaultrZ read_storagerrr=r>itemsizerr untyped_storage)r.rrrrrrgrrrrr<s"       zInputReader.storage)storage_offsetrhris_leafc Kst||d}t|}t|}t|}t|}tjg||j|d} t| ||||Wdn1s7wY|sut | j tj d} Wdn1sTwYt| ||||Wdn1spwYtj j| |ksJtj| ||j| | S)Nr)rhrr)Z memory_format)r_storage_offset_or_defaultr_is_leaf_or_default_requires_grad_or_defaultrdrrZno_gradset_Z enable_gradcloneZpreserve_format _subclasses meta_utils safe_is_leaf_utilsZset_tensor_metadatarappend) r.rrgrrrhrrmetadatatrrrrRs,      zInputReader.tensorcCs|j||Srx)rr)r.valrrrrss zInputReader.symintrxrrrrrr0s  !rc@s^eZdZddddZddZdddd efd d Zdd d ZddZdddZ dddZ dS) InputWriterF stable_hashcCs:g|_t|_||_|durt||dnd|_i|_dS)Nr)_lines itertoolsrstorage_counterrr r seen_storages)r.rrrrrr1s   zInputWriter.__init__cCs*dg}|dd|jD|d|S)Nzdef load_args(reader):css|]}d|VqdS)rXNr)rlrrrr{sz$InputWriter.lines..zload_args._version = 0)extendrr)r.rrrrliness  zInputWriter.linesNr device_hintrc Cst|}|j|}|dur|Sdt|j}d}tdt|kr'd|}d}|j}|jdkr9|dus7J|}td|krDd|}| } d} |j dur[|jjdkr[|j |} |j |d| d| ||d||j|<|S) Nbufrz , dtype_hint=metaz , device=z = reader.storage(, r[)r rrr]rrrrOrrrZ write_storagerr) r.rrr wsvZmaybe_dtype_hintrqrrrrrrrs0        zInputWriter.storagec CsNddlm}m}|j||j|jd}g}||td|jd| s/| t t | t d|jkr?| d|j|td|ksS| d|tj|}|rg|dd|Dtd|jkrw| d |jtjj|}td|kr| d ||j d d |t t |jg|d |dS)Nr)statically_known_truesym_eqr rzdtype=zstorage_offset=css"|] \}}|d|VqdS)=Nr)rkrrrrr{s z%InputWriter.tensor..zrequires_grad=zis_leaf=zreader.tensor(r) # )Z%torch.fx.experimental.symbolic_shapesrrrrrhrrrgrrstrtuplerrrrdrZget_tensor_metadatarr`rrrrrrrr2) r.rrrrrrZtensor_metadatarrrrrs<  zInputWriter.tensorcCs|jd|dt|t|ttfrV|jdt|D].\}}|d|d}t|tjr8| ||qt|t tj frG| ||q| ||q|jddSdS)Nruz# was unsupported type for dumping: z"""[])rrrOrrfr enumeraterdrrintSymIntr unsupported)r.rargrzaZname_irrrrs  zInputWriter.unsupportedcCs|jd|d|ddS)Nz reader.const(rz!, filtered out during compilation)rr)r.rrrrconstszInputWriter.constcCs0t|tjr |jj}|jd|d|dS)Nzreader.symint(r)rrdrrhintrr)r.rrrrrrs zInputWriter.symintr) r@rArBr1r rrrrr rrrrrrs   rryfuncr sym_shapesdefault_sym_shapecs~ddlm}dd|D}d|}t|}d|d}d|d } d } Gd d d } i} p5ifd ddtffdd } |j}|D]:\}}|dkrWqNt | |}|rw| \}}t | d}||}| ||| |<t | |}|r|d| |<qNdt|jvr| }|| d<t ||D]}| \}}}}t | d}||}t||| ||q| 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_abbrscSsi|]\}}||qSrrrrrr sz*aot_graph_input_parser..|z(_tensor_constant\d+): \"(z0)\[\s*(.*?)\s*\]\" = self\.(_tensor_constant\d+)(z)\[\s*(.*?)\s*\]z Sym\((s\d+)\)c@seZdZdZdS)z/aot_graph_input_parser..TensorContainerz#Container for tensors as attributesN)r@rArB__doc__rrrrTensorContainersr*cs,tvp dufddS)Ncs dS)Nz; not in symbolic_shapes and default sym shape not passed inrrrrrr's z=aot_graph_input_parser..get_sym_int..)rd_checkrr+)r$r#r+r get_sym_int$s   z+aot_graph_input_parser..get_sym_intrc sg}g}t|D]$\}}|}d|vr#|}||||q|r,|t|q|jr3tjntj}|||d}|D] } tj || q?|S)Nrw)rhr) rstriprrrerdZrandnZzerosrZ mark_dynamic) rgrhZresolved_shapeZ dynamic_dimsrzdimrw constructorrr)rr-rr gen_tensor+s   z*aot_graph_input_parser..gen_tensor,rr.)Ztorch.fx.graphr%r`r2valuesinspect getsourcer__annotations__researchgroupsrr&group signaturer^finditersetattr)r"rr#r$r%Z dtype_mapZ dtype_patternsourceZtensor_assignment_regexZ tensor_regexZsym_shape_regexr*rr1 annotationsrp annotationmatchZ data_typeZ shape_strrgrh containerZ attr_namerTr)r$rr-r#raot_graph_input_parsersF       rCr/csDttjtjfdd}fdd}t||S)z Decorator to cProfile a given function and save the result to disk on process exit. Args: filename: filename to save profile to cstfdd}|S)Ncs,z |i|WSwrx)enabledisable)rr)fnprofrrwrapperhsz3profile_to_file..decorator..wrapper) functoolswraps)rFrH)rG)rFr decoratorgsz"profile_to_file..decoratorc s.tjtddddS)Nz! Wrote profile to z+, view with: snakeviz z )Z dump_statssysstderrr;r4r5rr/rGrrsave_itrs z profile_to_file..save_it)cProfileZProfiler$r%r' expanduseratexitregister)r/rKrOrrNrprofile_to_file]s   rT)FF)F)ryNN)br)rRrrPrIrFr4rloggingr$r7r|rLrDr4 collectionsr importlibrtypingrrrrrdZtorch._prims_commonZ _prims_commonrZtorch._subclasses.meta_utilsrZtorch._dynamo.testingr r Z torch.multiprocessing.reductionsr Ztorch.utils._content_storer r rrrr getLoggerr@r=rZinductor_configZ is_fbcoderZlibfb.py.build_infoZlibfbr3Z extra_importspyZ build_infoZ BuildInfoZget_build_ruler*r6r2r<rrJrarL lru_cacherrrrrrrrrrrrrrrrZfloat32rrrrrrrrrrfrrrrCrTrrrrs         7 f   : )  Tx  c