Support wider range of types in FutureNCCL (#48502)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48502
This commit is part of a stack that reworks FutureNCCL in order to extract a generic CUDA-aware Future subclass. The stack deliberately breaks up this transition into elementary changes, to make it easier to verify that the behavior is preserved (or to highlight how it gets changed).
---
FutureNCCL restricted the values to be tensors, or (singleton) lists of tensors, or Python object that could be converted to either of those types. We need a CUDA future that can handle more generic types though.
The main challenge is extracting all DataPtrs from an arbitrary object. I think I found some ways of doing so, but I'd like some JIT experts to look into this and tell me if there are better ways. I'll add inline comments for where their input would be appreciated.
ghstack-source-id: 118180026
Test Plan: Unit tests (I should probably add new ones)
Reviewed By: wanchaol
Differential Revision: D25177562
fbshipit-source-id: 1ef18e67bf44543c70abb4ca152f1610dea4e533
