Add and adjust kernel launch checks (#46727)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46727
This adds kernel launch safety checks to a number of kernels. See D24309971 (https://github.com/pytorch/pytorch/commit/353e7f940f548e0a0cb3b420b4190b4624ae9b41) for context.
Test Plan: The existing pre-commit test rigs are used.
Reviewed By: ngimel
Differential Revision: D24334303
fbshipit-source-id: b6433f6be109fc8dbe789e91f3cbfbc31fd15951
