frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove some unnecessary functions from CUDAHooks (#59655)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59655

CUDAHooks is to be used solely when you need to call into CUDA
functionality from a context where you cannot directly link to
CUDA libraries.  Neither of hasPrimaryContext nor
getDevceIndexWithPrimaryContext (sic) needs to be used in such
contexts.  By moving them out of CUDAHooks and calling them
directly a dynamic dispatch can be skipped.

I also fixed the typo in getDev(i)ceIndexWithPrimaryContext

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

Test Plan: Imported from OSS

Reviewed By: ngimel

Differential Revision: D28972946

Pulled By: ezyang

fbshipit-source-id: edcd7a7b62aec97928f07fbf3bf413b9fb027517
This commit is contained in:
Edward Yang
2021-06-28 10:36:36 -07:00
committed by Facebook GitHub Bot
parent b52849b589
commit 85af24f52b
6 changed files with 37 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torch/csrc/cuda/Module.cpp
View File

@ -4,6 +4,7 @@
#include <ATen/CUDAGeneratorImpl.h>
#include <c10/cuda/CUDAFunctions.h>
#include <c10/cuda/CUDACachingAllocator.h>
#include <ATen/cuda/detail/CUDAHooks.h>
#ifdef USE_NCCL
#include <torch/csrc/cuda/python_nccl.h>
#endif
@ -285,7 +286,7 @@ PyObject * THCPModule_hasPrimaryContext(PyObject *_unused, PyObject *arg)
HANDLE_TH_ERRORS
THPUtils_assert(THPUtils_checkLong(arg), "invalid argument to has_primary_context");
int64_t device_index = static_cast<int64_t>(THPUtils_unpackLong(arg));
if (at::detail::getCUDAHooks().hasPrimaryContext(device_index)) {
if (at::cuda::detail::hasPrimaryContext(device_index)) {
Py_RETURN_TRUE;
} else {
Py_RETURN_FALSE;