Repository: pytorch
Issue: Divergent code is needed to record usage streams on different TensorImpl types
Labels: module: sparse,module: cuda,triaged,module: vmap

In https://github.com/pytorch/pytorch/pull/60230 I'm trying to add proper stream recording for tensors used across streams in streaming backward.

Naively writing
```C++
const auto guard = c10::impl::VirtualGuardImpl(c10::DeviceType::CUDA);
guard.recordDataPtrOnStream(var.storage().data_ptr(), stream); 
```
caused CI errors on some sparse and batched tensor tests, with messages like
[Cannot access storage of SparseTensorImpl](https://app.circleci.com/pipelines/github/pytorch/pytorch/338674/workflows/c8b6356e-d75f-4694-84b1-a65ac85f4e16/jobs/14233621)
[Cannot access storage of BatchedTensorImpl](https://app.circleci.com/pipelines/github/pytorch/pytorch/338674/workflows/c8b6356e-d75f-4694-84b1-a65ac85f4e16/jobs/14233623)

The TensorImpl children implementing these tensor types don't have a single storage() you can access. For example, sparse tensors tend to be backed by multiple data tensors (e.g. indices and values for COO aka SparseTensorImpl). So lack of a storage() for these is fair.

But to allow stream recording for different TensorImpls to fix CI, I had to write this beauty in [torch/csrc/autograd/input_buffer.cpp](https://github.com/pytorch/pytorch/pull/60230/files#diff-41f9e0edd7726a345a8fe7b597303b30329c30b6dbeb95fe5aef40dfd3981b02R19-R59):
```C++
  // look what you made me do >.<
  // Divergent paths for per-Impl stream recording that leak implementation
  // details of the impls should not be needed here.
  // See
  // TODO: clean this up when
  void record_stream_maybe_sparse(Variable& var, c10::Stream& stream) {
    const auto guard = c10::impl::VirtualGuardImpl(c10::DeviceType::CUDA);

    if (C10_UNLIKELY(at::isBatchedTensor(var))) {
      auto* impl = at::maybeGetBatchedImpl(var);
      if (impl) {
        guard.recordDataPtrOnStream(impl->value().storage().data_ptr(), stream);
      } else {
        TORCH_INTERNAL_ASSERT(false, "Expected batched tensor");
      }
    } else {
      switch (var.layout()) {
        case c10::kSparseCsr:
          {
            auto* impl = at::sparse_csr::get_sparse_csr_impl(var);
            guard.recordDataPtrOnStream(impl->values().storage().data_ptr(), stream);
            guard.recordDataPtrOnStream(impl->crow_indices().storage().data_ptr(), stream);
            guard.recordDataPtrOnStream(impl->col_indices().storage().data_ptr(), stream);
            break;
          }
        case c10::kSparse:
          {
            auto* impl = at::sparse::get_sparse_impl(var);
            guard.recordDataPtrOnStream(impl->values().storage().data_ptr(), stream);
            guard.recordDataPtrOnStream(impl->indices().storage().data_ptr(), stream);
            break;
          }
        case c10::kStrided:
          guard.recordDataPtrOnStream(var.storage().data_ptr(), stream);
          break;
        default:
          TORCH_INTERNAL_ASSERT(false, "Unknown layout in record_stream_maybe_sparse");
      }
    }
  }
}  // anonymous namespace
```
Ofc this one-off logic at my point of use is terrible design:
- Stream recording isn't unique to my point of use, other places that record streams might need this logic too.
- There might be other TensorImpl children I missed.
- If any Impl is added in the future, I'll need a new path to handle it.
- If any Impl's internal data tensors change, my logic must be updated to reflect it.

To improve things:
either the generic stream-recording helper should live in a central location TensorImpl authors know they must modify when they change or add a TensorImpl child,
or we should add a generic `record_use_on_stream(stream)` method to Tensor itself, which the author of each TensorImpl author would have to implement.

cc @ngimel @nikitaved @pearu @cpuhrsch @IvanYashchuk