Go through all the node.target and validate their Tensor arguments are having the allowed dtypes.
( # noqa: C901 # pyre-fixme[14]
self, target: _Target, args: Tuple[_Argument, ...], kwargs: Dict[str, _Argument]
)
| 63 | return kernel_arg |
| 64 | |
| 65 | def call_function( # noqa: C901 # pyre-fixme[14] |
| 66 | self, target: _Target, args: Tuple[_Argument, ...], kwargs: Dict[str, _Argument] |
| 67 | ) -> Any: |
| 68 | """ |
| 69 | Go through all the node.target and validate their Tensor arguments are having the allowed dtypes. |
| 70 | """ |
| 71 | if not isinstance(target, EdgeOpOverload) or not isinstance( |
| 72 | target._schema, EdgeDialectFunctionSchema |
| 73 | ): |
| 74 | if isinstance(target, HigherOrderOperator): |
| 75 | raise RunHigherOrderOperatorError("Can't run delegate") |
| 76 | return super().call_function(target, args, kwargs) # pyre-fixme[6] |
| 77 | |
| 78 | # TODO(gasoonjia): Update Optional[torch.dtype] to a concrete class to support mixed dtypes in tensorlist. |
| 79 | tensor_arg_types: Dict[str, Optional[torch.dtype]] = {} |
| 80 | for i, schema_arg in enumerate(target._schema.arguments): |
| 81 | if ( |
| 82 | isinstance(schema_arg.type, torch.TensorType) |
| 83 | or schema_arg.type == torch.OptionalType.ofTensor() |
| 84 | ): |
| 85 | kernel_arg = self._get_kernel_arg(schema_arg, i, args, kwargs) |
| 86 | if not isinstance(kernel_arg, torch.Tensor): |
| 87 | continue |
| 88 | tensor_arg_types[schema_arg.name] = kernel_arg.dtype |
| 89 | elif schema_arg.type == torch.ListType.ofTensors(): |
| 90 | kernel_arg = self._get_kernel_arg(schema_arg, i, args, kwargs) |
| 91 | if not isinstance(kernel_arg, list) or not all( |
| 92 | isinstance(kernel_arg[i], torch.Tensor) |
| 93 | for i in range(len(kernel_arg)) |
| 94 | ): |
| 95 | continue |
| 96 | if len(kernel_arg): |
| 97 | tensor_arg_types[schema_arg.name] = kernel_arg[0].dtype |
| 98 | else: |
| 99 | # If kernel_arg is an empty list, treat its type as None. |
| 100 | # FunctionDtypeConstraint.validate will take None as any legal dtype. |
| 101 | tensor_arg_types[schema_arg.name] = None |
| 102 | |
| 103 | ret_index = 0 |
| 104 | kernel_rets = self.node.meta["val"] |
| 105 | ret_iter = iter( |
| 106 | kernel_rets if isinstance(kernel_rets, Sequence) else [kernel_rets] |
| 107 | ) |
| 108 | for schema_ret in target._schema.returns: |
| 109 | name = schema_ret.name if schema_ret.name else f"__ret_{ret_index}" |
| 110 | kernel_ret = next(ret_iter) |
| 111 | if isinstance(schema_ret.type, torch.TensorType): |
| 112 | if isinstance(kernel_ret, torch.Tensor): |
| 113 | tensor_arg_types[name] = kernel_ret.dtype |
| 114 | ret_index += 1 |
| 115 | # Exceptionally rarely (basically only backwards ops) you might see an OptionalTensor returned. |
| 116 | # The schema of these ops though is typically -> (Tensor, Tensor ...). So the actual type |
| 117 | # returned in cpp is empty/undefined tensor. There is no analogy to this in python so it |
| 118 | # gets crudely mapped to None. To properly fix this core pytorch would have to change the |
| 119 | # schema to (Tensor?, ...) which is just never going to happen. So we have to handle this case |
| 120 | # here in the verifier and in memory planning as well. |
| 121 | elif kernel_ret is None: |
| 122 | tensor_arg_types[name] = schema_ret.default_value |
nothing calls this directly
no test coverage detected