pytorch / pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration
https://pytorch.org
Other
83.56k stars 22.54k forks source link

Decouple some generic APIs with CUDA building #131881

Open Danielmic opened 3 months ago

Danielmic commented 3 months ago

šŸš€ The feature, motivation and pitch

Currently there are some API are generic but defined in torch/csrc/cuda/Module.cpp and only exposed when USE_CUDA=ON

static void registerCudaPluggableAllocator(PyObject* module) {
  auto m = py::handle(module).cast<py::module>();  
  //  ...
  m.def(
      "_set_storage_access_error_msg", [](const at::Tensor& t, std::string s) {
        t.unsafeGetTensorImpl()
            ->release_storage_and_set_meta_custom_data_ptr_error_msg_(s);
      });

  m.def("_storage_Use_Count", [](size_t storage_impl_ptr) {
    // NOLINTNEXTLINE(performance-no-int-to-ptr)
    c10::StorageImpl* storage_impl = (c10::StorageImpl*)storage_impl_ptr;
    return c10::raw::weak_intrusive_ptr::use_count(storage_impl);
  });

  m.def(
      "_tensors_data_ptrs_at_indices_equal",
      [](py::list& tensors, py::list& data_ptrs, py::list& indices) {
        for (size_t i = 0, end = indices.size(); i < end; ++i) {
          auto index = indices[i].cast<int64_t>();
          auto t = tensors[index].cast<at::Tensor>();
          auto data_ptr = data_ptrs[index].cast<int64_t>();
          if (reinterpret_cast<int64_t>(t.data_ptr()) != data_ptr) {
            return false;
          }
        }
        return true;
      });
  // ...
}

Alternatives

move then from torch::cuda::initModule(module) to torch::initModule()

Additional context

No response

### Tasks

cc @ptrblck @msaroufim

Danielmic commented 3 months ago

I draft a PR here https://github.com/pytorch/pytorch/pull/131882

janeyx99 commented 3 months ago

These look to be private APIs -- could you elaborate on your use case for these functions?

Danielmic commented 2 months ago

I am using the PrivateUse1 backend, and I need to duplicate the same code implementation here. Therefore, I would like to see if these APIs can be made public. @janeyx99