U zhe@sUddlZddlZddlZddlZddlZddlZddlmZmZm Z m Z m Z m Z ddl Z ddlmZddlmZddlmZddlmZddlmZmZddlZ ddlmZdd d d gZd d dZddddddhZd?eejeejdddZ iZ!ej"ed fe#d<Gdd d Z$ej%GdddZ&eedddZ'edddd Z(edd!d"d#Z)eej*eefd$d%d&Z+edd'd(d)Z,ej-e.d*d+d,Z/eejdd-d.d/Z0ej1edd0d1d2Z2d3d4Z3d5d6Z4d@d8d9Z5d:d;Z6dAd=d>Z7dS)BN)FunctionSchema OperatorName SchemaKindBaseTypeListTypeBaseTy)AbstractImplCtx)get_ctx)autograd_kernel_indirectionconstruct_autograd_kernel) infer_schema custom_opCustomOpr rCPUCUDA)cpucudaZprimZprimsZatenattorchZpytorch)qualname manual_schemareturncsfdd}|S)aR Creates a new CustomOp object. WARNING: if you're a user, please do not use this directly (instead use the torch._custom_ops APIs). Also please see the following for a detailed guide on custom ops. https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk In PyTorch, defining an op (short for "operator") is a two step-process: - we need to define (create) the op - we need to implement behavior for how the operator interacts with various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc. This entrypoint defines the CustomOp object (the first step); you must then perform the second step by calling various methods on the CustomOp object. This API is used as a decorator (see examples). Arguments: qualname (str): Should be a string that looks like "namespace::operator_name". Operators in PyTorch need a namespace to avoid name collisions; a given operator may only be created once. If you are writing a Python library, we recommend the namespace to be the name of your top-level module. The operator_name must be the same as the name of the function you pass to custom_op (see examples). manual_schema (Optional[str]): Each PyTorch operator needs a schema that tells PyTorch the types of the inputs/outputs. If None (default), we will infer the schema from the type annotations on the function (see examples). Otherwise, if you don't want to use type annotations, you may provide us the schema string. Example:: >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) >>> import numpy as np >>> from torch import Tensor >>> >>> # Step 1: define the CustomOp. >>> # We need to provide the decorator a "prototype function" >>> # (a function with Python ellipses as the body). >>> @custom_op("my_library::numpy_sin") >>> def numpy_sin(x: Tensor) -> Tensor: >>> ... >>> >>> # numpy_sin is now an instance of class CustomOp >>> print(type(numpy_sin)) >>> >>> # Step 2: Register an implementation for various PyTorch subsystems >>> >>> # Register an implementation for CPU tensors >>> @numpy_sin.impl('cpu') >>> def numpy_sin_impl_cpu(x): >>> return torch.from_numpy(np.sin(x.numpy())) >>> >>> # Register an implementation for CUDA tensors >>> @numpy_sin.impl('cuda') >>> def numpy_sin_impl_cuda(x): >>> return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device) >>> >>> x = torch.randn(3) >>> numpy_sin(x) # calls numpy_sin_impl_cpu >>> >>> x_cuda = x.cuda() >>> numpy_sin(x) # calls numpy_sin_impl_cuda c s(t|stdt|t\}}t||j|krXtdd|d|jddkrht|n}||}t |}t |dk rt ||t |d}||t||j}t|||||dd}|j|_|j|_|j|_t ||jd tt|tj|ttt||S) NzDcustom_op(...)(func): Expected `func` to be a Python function, got: zcustom_op(qualname='z-', ...)(func): expected `func` to have name 'z ' but got 'zX'. Please either change the name of `func` or the qualname that is passed to `custom_op`FRAGMENTT_private_accessAutograd)inspect isfunction ValueErrortypeparse_qualnamevalidate_namespace__name__r rparsevalidate_schema validate_function_matches_schemalibraryLibrarydefinefind_ophandle_or_thrownamer __module____doc__impl_opnamer weakrefproxyr_C#_dispatch_set_report_error_callback functoolspartialreport_error_callback) funcnsr+schema schema_strfunction_schemalibophandleresultrrG/var/www/html/venv/lib/python3.8/site-packages/torch/_custom_op/impl.pyinnervs>           zcustom_op..innerr@)rrrBr@r?rAr0sF+global_registrycseZdZdZddfdd ZddZd,d d Zd d Zd dZddZ ddZ ddZ d-e j ee jefe jdddZddZe jdddZd.e jdddZdd Zd!d"Zd#d$Zd%d&Zd/d'd(Zd0d*d+ZZS)1ra3Class for custom operators in PyTorch. Use the CustomOp API to create user-defined custom operators that behave just like regular PyTorch operators (e.g. torch.sin, torch.mm) when it comes to various PyTorch subsystems (like torch.compile). To construct a `CustomOp`, use `custom_op`. Frcsht|std|d|}||_||_||_||_||_||_d|_ i|_ d|_ |t |j<dS)Nz|The CustomOp constructor is private and we do not guarantee BC for it. Please use custom_op(...) to create a CustomOp object::F) super__init__ RuntimeError_schema_cpp_ns_lib _ophandler/ _qualnamer#_impls'_registered_autograd_kernel_indirectionrC)selfr<cpp_nsr9 operator_namer=rr+ __class__r@rArFs  zCustomOp.__init__cCs0|jr t|j|jtt|dd|_dS)NrT)rNAssertionErrorrJr.r/r r0r1rOr@r@rA%_register_autograd_kernel_indirections z.CustomOp._register_autograd_kernel_indirectionc Cs||rJ|j|}|dk s t|j}td|d|jd|d|d tt |}|j d|j }t |||j|<dS)NzAttempting to register a z impl for operator z that already has a z impl registered from Python at z. This is not supported.:) _has_implrMrTlocationrGrLr getframeinfosys _getframefilenamelinenoFuncAndLocation)rOkindr7 stacklevelZfunc_and_locationrZframer@r@rA_register_impls   zCustomOp._register_implcCs |j|SNrMrOrar@r@rA _get_implszCustomOp._get_implcCs ||jkSrerfrgr@r@rArYszCustomOp._has_implcCs6|`ttj|j}t||jr*t||jt|j =dSre) rJgetattrropsrIhasattrr/delattrrCrL)rO opnamespacer@r@rA_destroys   zCustomOp._destroycCsd|jdS)Nz)rLrUr@r@rA__repr__szCustomOp.__repr__cOstj|jf||}|Sre)r2Z_dispatch_call_boxedrK)rOargskwargsr>r@r@rA__call__szCustomOp.__call__) device_typesrcs6ttrgD] }t|qfdd}|S)aRegister an implementation for a device type for this CustomOp object. WARNING: if you're a user, please do not use this directly (instead use the torch._custom_ops APIs). Also please see the following for a detailed guide on custom ops. https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk If the CustomOp is passed multiple Tensor inputs with different device types, it will dispatch to the registered implementation for the highest priority device type among those present. The supported device types, in order of priority, are {'cuda', 'cpu'}. This API is used as a decorator (see examples). Arguments: device_types (str or Iterable[str]): the device type(s) to register the function for. Examples:: >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) >>> import numpy as np >>> from torch import Tensor >>> >>> @custom_op("my_library::numpy_cos") >>> def numpy_cos(x: Tensor) -> Tensor: >>> ... >>> >>> # Register an implementation for CPU Tensors >>> @numpy_cos.impl('cpu') >>> def numpy_cos_impl_cpu(x): >>> return torch.from_numpy(np.cos(x.numpy())) >>> >>> # Register an implementation for CUDA Tensors >>> @numpy_cos.impl('cuda') >>> def numpy_cos_impl_cuda(x): >>> return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device) >>> >>> x = torch.randn(3) >>> numpy_cos(x) # calls numpy_cos_impl_cpu >>> >>> x_cuda = x.cuda() >>> numpy_cos(x) # calls numpy_cos_impl_cuda csJtD]<}|j||dt|}tjj||q|S)Nrb)set_check_doesnt_have_library_implrdSUPPORTED_DEVICE_TYPE_TO_KEYr'r.rJr/)f device_typeZ dispatch_key _stacklevelrsrOr@rArB8s   zCustomOp.impl..inner) isinstancestrvalidate_device_type)rOrsr{ryrBr@rzrAr.s .  z CustomOp.implcCs@||rdSt|}t|j|r.innerr@)rOrBr@rUrA impl_factoryLs zCustomOp.impl_factorycsfdd}|S)aj Register an abstract implementation for this operator. WARNING: please do not use this directly (and instead use the torch._custom_ops APIs). Also please see the following for a detailed guide on custom ops. https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk An "abstract implementation" specifies the behavior of this operator on Tensors that carry no data. Given some input Tensors with certain properties (sizes/strides/storage_offset/device), it specifies what the properties of the output Tensors are. The abstract implementation has the same signature as the operator. It is run for both FakeTensors and meta tensors. To write an abstract implementation, assume that all Tensor inputs to the operator are regular CPU/CUDA/Meta tensors, but they do not have storage, and you are trying to return regular CPU/CUDA/Meta tensor(s) as output. The abstract implementation must consist of only PyTorch operations (and may not directly access the storage or data of any input or intermediate Tensors). This API is used as a decorator (see examples). Examples:: >>> import numpy as np >>> from torch import Tensor >>> >>> # Example 1: an operator without data-dependent output shape >>> @custom_op('my_library::custom_linear') >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor: >>> ... >>> >>> @custom_linear.impl_abstract() >>> def custom_linear_abstract(x, weight): >>> assert x.dim() == 2 >>> assert weight.dim() == 2 >>> assert bias.dim() == 1 >>> assert x.shape[1] == weight.shape[1] >>> assert weight.shape[0] == bias.shape[0] >>> assert x.device == weight.device >>> >>> return (x @ weight.t()) + bias >>> >>> # Example 2: an operator with data-dependent output shape >>> @custom_op('my_library::custom_nonzero') >>> def custom_nonzero(x: Tensor) -> Tensor: >>> ... >>> >>> @custom_nonzero.impl_abstract() >>> def custom_nonzero_abstract(x): >>> # Number of nonzero-elements is data-dependent. >>> # Since we cannot peek at the data in an abstract impl, >>> # we use the ctx object to construct a new symint that >>> # represents the data-dependent size. >>> ctx = torch._custom_op.get_ctx() >>> nnz = ctx.create_unbacked_symint() >>> shape = [x.dim(), nnz] >>> result = x.new_empty(shape, dtype=torch.long) >>> return result >>> >>> @custom_nonzero.impl(['cpu', 'cuda']) >>> def custom_nonzero_impl(x): >>> x_np = to_numpy(x) >>> res = np.stack(np.nonzero(x_np), axis=1) >>> # unbacked symbolic ints in PyTorch must be >= 2, so we >>> # constrain the range to at least 2 >>> if res.shape[0] <= 1: >>> raise RuntimeError("not supported") >>> return torch.tensor(res, device=x.device) csZjdddjjtfdd}jj |dS)Nabstractrtc s>fdd}tjj|||W5QRSQRXdS)NcstdddS)Nz.inner..f_with_ctx..error_on_ctx)r_libraryZ abstract_implZset_ctx_getter)rprqrrxrZrr@rA f_with_ctxs z9CustomOp.impl_abstract..inner..f_with_ctxMeta) $_check_doesnt_have_library_meta_implrdrhrZrLr4wrapsrJr.r/)rxrr{rOrrArBs z%CustomOp.impl_abstract..innerr@rOr{rBr@rrA impl_abstractVsHzCustomOp.impl_abstractc sfdd}jtjkr*|dj}js>|dt|dksNttdd|D}|rl|dtt j d tt j d tt j d tt j d tt jd ttt jddi}jD]0}|j|krq|dt|d|jdqdS)Ncs td|djddS)NzCCannot use torch._custom_ops APIs to register backward formula for z. Got operator z with schema: )rGrL)detailr9rOr@rAerrorsz4CustomOp._check_can_register_backward..errorznon-functional operatorzoperator with no returnsrcss"|]}|jdk o|jj VqdSre) annotationZis_write).0rr@r@rA sz8CustomOp._check_can_register_backward..zoperator that returns viewsintSymIntboolfloatTensorz List[Tensor]zoperator with return not in z (got ))rHrarZ functionalreturnslenrTanyrrrrrrrrr listvalues)rOrZretsZis_non_mutating_viewZallowed_return_typesretr@rrA_check_can_register_backwards<  z%CustomOp._check_can_register_backwardcCs^|jr dSt|jdr*td|jddD]*}t|j|r.td|jd|dq.dS)NCompositeImplicitAutogradz3impl_backward/impl_save_for_backward: the operator a3 already has an implementation for this device type via a pre-existing registration to DispatchKey::CompositeImplicitAutograd.CompositeImplicitAutograd operators do not need an autograd formula; instead, the operator will decompose into its constituents and those can have autograd formulas defined on them.)rZ AutogradCPUZ AutogradCUDAz; already has an Autograd kernel registered to DispatchKey::z vi a pre-existing torch.library or TORCH_LIBRARY registration. Please either remove those registrations or don't use the torch._custom_ops APIs)rNr2%_dispatch_has_kernel_for_dispatch_keyrLrG)rOrr@r@rA(_check_doesnt_have_library_autograd_impls  z1CustomOp._check_doesnt_have_library_autograd_implcCsr|drdSt|jdr.t|jds.dSt|jdrNtd|jdt|jdrntd|jddS)NrZCompositeExplicitAutogradrrz!impl_abstract(...): the operator a- already has an implementation for this device type via a pre-existing registration to DispatchKey::CompositeImplicitAutograd.CompositeImplicitAutograd operators do not need an abstract impl; instead, the operator will decompose into its constituents and those can have abstract impls defined on them.z already has an DispatchKey::Meta implementation via a pre-existing torch.library or TORCH_LIBRARY registration. Please either remove that registration or don't call impl_abstract.)rYr2rrLrGrUr@r@rArs     z-CustomOp._check_doesnt_have_library_meta_implc CsX|dst|dstt|j|j|t|j|dj|dj}| d|dS)Nbackwardsave_for_backwardautograd) rYrTr rH_output_differentiabilityget_oprLrhr7rd)rOZkernelr@r@rA_register_autograd_kernel"s  z"CustomOp._register_autograd_kernelcsfdd}|S)zyRegister a function that tells us what to save for backward. Please see impl_backward for more details. csDjsjd|ddr@dS)Nrrtr)rrrNrVrdrYrrrr@rArB3s z.CustomOp.impl_save_for_backward..innerr@rr@rrAimpl_save_for_backward.szCustomOp.impl_save_for_backwardNcsldk rXfdd}tts$|D]}t|ts(|q(tjjtkrX|fdd}|S)a2Registers a backward formula. WARNING: if you're a user, please do not use this directly (instead use the torch._custom_ops APIs). Also please see the following for a detailed guide on custom ops. https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk In order for the CustomOp to work with autograd, you need to register a backward formula. There are two pieces to this: 1. You must give us a function to specify what to save for backward. Call this the "save for backward" function. 2. You must give us a function that computes gradients. Call this the "backward" function. Use `impl_save_for_backward` to define a "save for backward" function that specifies what gets saved for backward. The function should accept two arguments ``(inputs, output)`` and return the quantities to be saved for backward. During runtime, when you call the CustomOp, PyTorch will invoke the "save for backward" function with the inputs and output of the CustomOp. Use `impl_backward` to define the "backward" function. The backward function must accept ``(ctx, saved, *grads)``: - ``ctx`` is a context object where we may provide information - ``saved`` is exactly what gets returned from the "save for backward" function - ``grads`` is one or more gradients. The number of gradients matches the number of outputs of the CustomOp. The backward function must return a dict that maps the name of an input to the CustomOp to its corresponding gradient. All inputs that were declared to be Tensors in the CustomOp definition must be accounted for in the dict. The gradient may be a Tensor or None. NcstddS)Nzimpl_backward(output_differentiability): expected output_differentiability to be a list of bools with length equal to the number of outputs of this CustomOp got: rr@)output_differentiabilityr@rAyellcsz$CustomOp.impl_backward..yellcsJjsjd|d_drFdS)Nrrtr)rrrNrVrdrrYrrr{rrOr@rArBrs z%CustomOp.impl_backward..inner)r|rrrrHr)rOrr{rdiffrBr@rrA impl_backward=s%    zCustomOp.impl_backward)rW)rW)rW)rW)NrW)r#r, __qualname__r-rFrVrdrhrYrnrorrtypingUnionr}IterableCallabler.rvrrrrrrrr __classcell__r@r@rRrArs.   =  f%) c@s eZdZUejed<eed<dS)r`r7rZN)r#r,rrr__annotations__r}r@r@r@rAr`~s  r`)rPrQcCs0|jdkrdn|j}t|dt|j|S)NrD) overload_namer2Z_dispatch_find_schema_or_throwr}r+)rPrQrr@r@rAr*s r*)r8rcCs:d|krtd|d|tkr6td|d|ddS)N.zcustom_op(..., ns="zC"): expected ns to not contain any . (and be a valid variable name)zcustom_op(..., ns='z'): 'z9' is a reserved namespace, please choose something else. )r RESERVED_NS)r8r@r@rAr"s r")r9rcCs:tjj|std||jjdk r6td|dS)Nzcustom_op only supports functional operators (ops that do not mutate any inputs, do not return views of the inputs, and has at least one return). Got the following non-functional schema: zUcustom_op does not support arguments named 'self'. Please rename your argument. Got: )rrutilsZis_functional_schemar argumentsZself_arg)r9r@r@rAr%s r%)rrcCsR|dd}t|dkr(td|dd|dkrBtd||d|dfS) NrDr rWz$Expected there to be a namespace in z;, i.e. The operator name should look something like ns::foorzThe torch.custom_ops APIs do not handle overloads, i.e. operator names with '.' in them. Please name your operator something like ns::foo. Got: r)splitrr)rnamesr@r@rAr!s    r!)ryrcCs&|tkr"td|dtddS)NzCustomOp.impl(device_types=[z(, ...]): we only support device_type in r)rwrkeys)ryr@r@rAr~sr~)paramrcCs|jtjjtjjfkSre)rar ParameterPOSITIONAL_OR_KEYWORD KEYWORD_ONLY)rr@r@rAsupported_paramsr)r9r7rcst|tddjDs0tdtddjDsVjtjj k rdtdddjD}ddjD}fd d fd d fd d}||j j ||j j dS)Ncss|]\}}t|VqdSre)rr_pr@r@rArsz3validate_function_matches_schema..zcustom_op(..., manual_schema)(func): positional-only args, varargs, and kwargs are not supported. Please rewrite `func` to not have them. Got `func` with signature: css |]\}}|jtjjk VqdSre)rrremptyrr@r@rArszcustom_op(..., manual_schema)(func): When passing in a manual schema, we expect `func` to have no type annotations to avoid ambiguity. Got `func` with signature: cSs&g|]\}}|jtjjkr||fqSr@)rarrrrr+rr@r@rA sz4validate_function_matches_schema..cSs&g|]\}}|jtjjkr||fqSr@)rarrrrr@r@rArscstdddS)Nzcustom_op(..., manual_schema)(func): When passing in a manual schema, we expect `func`'s signature to match `manual_schema` (aside from type annotations). func's signature: , manual_schema: rr@r9sigr@rArsz/validate_function_matches_schema..errorcstdddS)Nzycustom_op(..., manual_schema)(func): neither func nor manual_schema should have default arguments. Got func's signature: rrr@rr@rAerror_default_argssz.error_default_argscs`t|t|krt||D]:\\}}}||jkr<|jtjjk sT|jdk r q dSre)rzipr+defaultrrr)Zsig_argsZ schema_argsr+rarg)rrr@rAcompares z1validate_function_matches_schema..compare) r signatureall parametersitemsrrreturn_annotation SignaturerrZflat_positionalZflat_kwarg_only)r9r7 positionalZ kwargonlyrr@)rrr9rrAr&s2   r&)rrrc Cst|dkrt|d|dkr,t|d|dkr\|}t|d|d|d|d t|d |d dS) NZ Undefineda: There were no Tensor inputs to this operator (e.g. you passed an empty list of Tensors). If your operator is a factory function (that is, it takes no Tensors and constructs a new one), then please use CustomOp.impl_factory to register an implementation for itrz: when running with device='Meta' tensors: there is no abstract impl registered for this CustomOp. Please register one via CustomOp.impl_abstract to get this CustomOp to work with Meta tensors)rrz: when running with device='z' tensors: there is no zW impl registered for this CustomOp. Please register one via CustomOp.impl(device_type='z')z%: No implementation for dispatch key z. It is likely that we have not added this functionality yet, please either open an issue or if you're feeling adventurous, use the low-level torch.library API)NotImplementedErrorlower)rrZdevicer@r@rAr6s r6cCs\|j}tj|d}|dd}t|j}|dd}t |}t |||||ddS)NrrDTr) namespacerr'r(r+rr}rHrr$r)opr8r<r+r:r9r@r@rAcustom_op_from_existing%s  rcsffdd}t\}}ttj|s*|ttj|}t||sF|t||}t|ds`||jS)NcstdddS)NzCould not find the operator z~. Please make sure you have already registered the operator and (if registered from C++) loaded it via torch.ops.load_library.rr@rr@rAerror_not_found1s zget_op..error_not_foundr)r!rkrrjrir)rrr8r+rmpacketr@rrAr0s       rFcCs8|tkrt|S|s$td|dt|}t|}|S)NzCould not find custom op "z5". Did you register it via the torch._custom_ops API?)rCrGrr)rZalso_check_torch_libraryoverloadr>r@r@rA_find_custom_opCs rcCsF|tjjjkrdStjjj|}|dkr,dS|ds:dS|djS)Nr)rZ _custom_opr.rCrYrhr7)rrr@r@rAget_abstract_implOs rTc Cs|d\}}||}t|}t||rr@r@rA_custom_op_with_schemaZs    r)N)F)T)8 dataclassesr4rr\rr0Ztorchgen.modelrrrrrrrZtorch._Cr2Z torch.libraryr'Ztorch._library.abstract_implrr rr r Ztorch._library.infer_schemar __all__rwrr}OptionalrrrCDictrr dataclassr`r*r"r%Tupler!r~rrrr&Anyr6rrrrrr@r@r@rAsn        zT    A