[PyTorch Mobile] Generate Kernel dtype selection code in selected_mobile_ops.h during the build (#49279)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49279
Now that the YAML files for tracing based selective build optionally have the information regarding the selected kernel function dtypes, we can start generating constexpr selection code in the include file (`selected_mobile_ops.h`) to make the inclusion of code for specific dtypes selective based on compile time decisions.
The way this is done is that if we detect that the code for a specific dtype should not be in the binary, we add an exception (throw) statement just before the method is called (see the first diff in this dtack) and allow the compiler to optimize away the rest of the function's body. This has the advantage of allowing the compiler to know the lambda's return type (since it's inferred from the `return` statements in the body of the method, and if we compile out all the cases, then the compiler won't know the return type and it will result in a compilation error).
The generated `<ATen/selected_mobile_ops.h>` is being used (included) in `Dispatch.h`. In case `XPLAT_MOBILE_BUILD` is not defined, then we should include code for all kernel dtypes (non-selective build).
When merging, we need to handle the case of both older and newer (tracing based) operator lists. If we detect any operator that includes all overloads, it indicates that an old style operator list is part of the build, and we need to `include_all_kernel_dtypes` for this build.
ghstack-source-id: 119439497
Test Plan:
For Segmentation v220, here is one of the intermediate generated YAML files (selected_operators.yaml): {P154480509}
and here is the generated `selected_mobile_ops.h` file: {P159808798}
Here is the `selected_mobile_ops.h` file for lite_predictor (which includes all ops and all dtypes): {P159806443}
Continuous build for ~8 checked-in models validates that the selection code works as expected when we build based on dtype selection.
Reviewed By: iseeyuan
Differential Revision: D25388949
fbshipit-source-id: 1c182a4831a7f94f7b152f02dbd3bc01c0d22443
