mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
b8dfb45ac2
Summary: This PR is to update PyTorch with the following cub changes: - Starting cub 1.13.1, cub requires users to define `CUB_NS_QUALIFIER` if `CUB_NS_PREFIX` is also defined. Besides that, a new mechanism `CUB_WRAPPED_NAMESPACE` is added. And I do the following change to PyTorch: - Starting CUDA 11.5, define `CUB_WRAPPED_NAMESPACE` globally as an nvcc flag. - Fix caffe2 failures caused by the above change. - Add a `aten/src/ATen/cuda/cub_definitions.cuh` that defines helper macros about feature availability. Pull Request resolved: https://github.com/pytorch/pytorch/pull/66219 Reviewed By: bdhirsh Differential Revision: D31626931 Pulled By: ngimel fbshipit-source-id: 97ebf5ef671ade8bf46d0860edc317f22660f26d
18 lines
529 B
Plaintext
18 lines
529 B
Plaintext
#pragma once
|
|
|
|
// cub sort support for CUB_WRAPPED_NAMESPACE is added to cub 1.13.1 in:
|
|
// https://github.com/NVIDIA/cub/pull/326
|
|
// CUB_WRAPPED_NAMESPACE is defined globally in cmake/Dependencies.cmake
|
|
// starting from CUDA 11.5
|
|
#if defined(CUB_WRAPPED_NAMESPACE) || defined(THRUST_CUB_WRAPPED_NAMESPACE)
|
|
#define USE_GLOBAL_CUB_WRAPPED_NAMESPACE() true
|
|
#else
|
|
#define USE_GLOBAL_CUB_WRAPPED_NAMESPACE() false
|
|
#endif
|
|
|
|
#if USE_GLOBAL_CUB_WRAPPED_NAMESPACE()
|
|
namespace caffe2 {
|
|
namespace cub = ::CUB_WRAPPED_NAMESPACE::cub;
|
|
}
|
|
#endif
|