Add boxed fallback (#26368)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/26368
A boxed fallback function is a per-TensorTypeId generic function that is called if there is no operator specific function (either for the TensorTypeId, or the UndefinedTensorId fallback for that operator).
The outline of the patch:
* Add a new `boxed_fallback_table_` in the top-level ATenDispatch struct (this is NOT per operator) for storing boxed fallback functions. Boxed fallback functions provisionally have the interface `void(const char* schema, torch::jit::Stack*)`; we expect to replace the `const char* schema` with something more detailed in the near future.
* Boxing and unboxing is not supported for all arguments, as IValue doesn't cover the full set of types that are in ATen. The `supports_boxed_fallback` type trait tests if boxing is possible. The list of exclusions here was experimentally generated. One notable exclusion is that we cannot handle any mutable functions, as they return a `Tensor&` which we cannot conveniently manufacture from a Tensor on an IValue stack.
* The actual boxed fallback is handled in two phases. First, we do a (compile time) test to see if boxing/unboxing is supported. If it is, we do a runtime test to check if a fallback function is installed. If both conditions are met, we allocate a stack, push our arguments onto it, and then call the boxed fallback function. The return value is expected to be a single argument on the top of the stack; we retrieve it, unpack it, and return it to the user.
At present, there are no tests for this diff.
This diff also makes `multi_dispatch_tensor_type_set` safer by explicitly specifying that it takes all arguments as references. This prevents the function from accidentally moving in rvalue references (this never actually happened because none of the overloads of apply move out their arguments; but they could have.)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Differential Revision: D17448555
Test Plan: Imported from OSS
Pulled By: ezyang
fbshipit-source-id: 43a5be7bdfb5c0b466d01a55380cb79e6d938044