libtorch.zip

#include <c10d/ProcessGroupGloo.hpp> #include <netdb.h> #include <sys/socket.h> #include <sys/types.h> #include <unistd.h> #include <type_traits> #include <gloo/allgather.h> #include <gloo/allgatherv.h> #include <gloo/allreduce.h> #include <gloo/barrier.h> #include <gloo/broadcast.h> #include <gloo/gather.h> #include <gloo/reduce.h> #include <gloo/scatter.h> #include <ATen/SparseTensorUtils.h> #ifdef USE_CUDA #include <ATen/cuda/CUDAEvent.h> #include <ATen/cuda/Exceptions.h> #include <ATen/cuda/PinnedMemoryAllocator.h> #include <c10/cuda/CUDACachingAllocator.h> #include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAStream.h> #endif #include <gloo/config.h> #include <gloo/rendezvous/context.h> #include <gloo/rendezvous/prefix_store.h> #if GLOO_HAVE_TRANSPORT_TCP #include <gloo/transport/tcp/device.h> #endif #if GLOO_HAVE_TRANSPORT_UV #include <gloo/transport/uv/device.h> #endif // On Linux, check that the tcp transport is available. #ifdef __linux__ #if !GLOO_HAVE_TRANSPORT_TCP #error "Expected the tcp transport to be available on Linux." #endif #endif // On macOS, check that the uv transport is available. #ifdef __APPLE__ #if !GLOO_HAVE_TRANSPORT_UV #error "Expected the uv transport to be available on macOS." #endif #endif #define GENERATE_ALL_TYPES(type, func, args...) \ switch (type) { \ case ::at::ScalarType::Float: \ func<float>(args); \ break; \ case ::at::ScalarType::Double: \ func<double>(args); \ break; \ case ::at::ScalarType::Half: \ func<gloo::float16>(args); \ break; \ case ::at::ScalarType::Char: \ func<int8_t>(args); \ break; \ case ::at::ScalarType::Byte: \ func<uint8_t>(args); \ break; \ case ::at::ScalarType::Int: \ func<int32_t>(args); \ break; \ case ::at::ScalarType::Long: \ func<int64_t>(args); \ break; \ default: \ throw std::runtime_error("Invalid scalar type"); \ } namespace c10d { namespace { // Wrap c10d store as Gloo store class GlooStore : public ::gloo::rendezvous::Store { public: GlooStore(const std::shared_ptr<::c10d::Store>& store) : store_(store) {} void set(const std::string& key, const std::vector<char>& value) override { std::vector<uint8_t> tmp(value.begin(), value.end()); store_->set(key, tmp); } std::vector<char> get(const std::string& key) override { auto value = store_->get(key); return std::vector<char>(value.begin(), value.end()); } void wait(const std::vector<std::string>& keys) override { store_->wait(keys, Store::kDefaultTimeout); } void wait( const std::vector<std::string>& keys, const std::chrono::milliseconds& timeout) override { store_->wait(keys, timeout); } protected: std::shared_ptr<::c10d::Store> store_; }; typedef void (*ReduceFunc)(void*, const void*, const void*, size_t); template < typename T, typename std::enable_if<!std::is_integral<T>::value, int>::type = 0> ReduceFunc toFunction(const ReduceOp& r) { switch (r) { case ReduceOp::SUM: return ReduceFunc(&::gloo::sum<T>); case ReduceOp::PRODUCT: return ReduceFunc(&::gloo::product<T>); case ReduceOp::MIN: return ReduceFunc(&::gloo::min<T>); case ReduceOp::MAX: return ReduceFunc(&::gloo::max<T>); case ReduceOp::BAND: throw std::runtime_error( "Cannot use ReduceOp.BAND with non-integral dtype"); break; case ReduceOp::BOR: throw std::runtime_error( "Cannot use ReduceOp.BOR with non-integral dtype"); break; case ReduceOp::BXOR: throw std::runtime_error( "Cannot use ReduceOp.BXOR with non-integral dtype"); break; case ReduceOp::UNUSED: break; } throw std::runtime_error("Unhandled ReduceOp"); } // Bitwise AND with SFINAE guard for integral types. template < typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0> void band(void* c, const void* a, const void* b, size_t n) { auto tc = static_cast<T*>(c); auto ta = static_cast<const T*>(a); auto tb = static_cast<const T*>(b); for (size_t i = 0; i < n; i++) { tc[i] = ta[i] & tb[i]; } } // Bitwise OR with SFINAE guard for integral types. template < typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0> void bor(void* c, const void* a, const void* b, size_t n) { auto tc = static_cast<T*>(c); auto ta = static_cast<const T*>(a); auto tb = static_cast<const T*>(b); for (size_t i = 0; i < n; i++) { tc[i] = ta[i] | tb[i]; } } // Bitwise XOR with SFINAE guard for integral types. template < typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0> void bxor(void* c, const void* a, const void* b, size_t n) { auto tc = static_cast<T*>(c); auto ta = static_cast<const T*>(a); auto tb = static_cast<const T*>(b); for (size_t i = 0; i < n; i++) { tc[i] = ta[i] ^ tb[i]; } } template < typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0> ReduceFunc toFunction(const ReduceOp& r) { switch (r) { case ReduceOp::SUM: return ReduceFunc(&::gloo::sum<T>); case ReduceOp::PRODUCT: return ReduceFunc(&::gloo::product<T>); case ReduceOp::MIN: return ReduceFunc(&::gloo::min<T>); case ReduceOp::MAX: return ReduceFunc(&::gloo::max<T>); case ReduceOp::BAND: return ReduceFunc(&band<T>); case ReduceOp::BOR: return ReduceFunc(&bor<T>); case ReduceOp::BXOR: return ReduceFunc(&bxor<T>); case ReduceOp::UNUSED: break; } throw std::runtime_error("Unhandled ReduceOp"); } template <typename T, typename O> void setInputs(O& opts, std::vector<at::Tensor>& tensors) { opts.setInputs(getDataPointers<T>(tensors), tensors[0].numel()); } template <typename T, typename O> void setInput(O& opts, at::Tensor& tensor) { opts.setInput(getDataPointer<T>(tensor), tensor.numel()); } template <typename T, typename O> void setOutputs(O& opts, std::vector<at::Tensor>& tensors) { opts.setOutputs(getDataPointers<T>(tensors), tensors[0].numel()); } template <typename T, typename O> void setOutput(O& opts, at::Tensor& tensor) { opts.setOutput(getDataPointer<T>(tensor), tensor.numel()); } template <typename T, typename O> void setOutput(O& opts, at::Tensor& tensor, std::vector<size_t>& counts) { opts.setOutput(getDataPointer<T>(tensor), counts); } #ifdef USE_CUDA at::Tensor pinnedLike(at::Tensor& tensor) { auto* allocator = at::cuda::getPinnedMemoryAllocator(); auto storage = c10::Storage( tensor.dtype(), at::detail::computeStorageSize(tensor.sizes(), tensor.strides()), allocator, /*resizable=*/false); return at::empty({0}, tensor.options().device(at::kCPU)) .set_(storage, 0, tensor.sizes(), tensor.strides()); } // This function initializes a vector of CUDA streams, one for every // tensor in the input tensor vector, and ensures that these streams are // synchr