FM / cuda

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/cuda/

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/cuda

Name	Size	Mode	Actions
comm.h	1525	0644	edit dl rm
device_set.h	152	0644	edit dl rm
Event.h	487	0644	edit dl rm
Module.h	503	0644	edit dl rm
nccl.h	5037	0644	edit dl rm
override_macros.h	1887	0644	edit dl rm
python_comm.h	114	0644	edit dl rm
python_nccl.h	616	0644	edit dl rm
restore_macros.h	820	0644	edit dl rm
serialization.h	524	0644	edit dl rm
Storage.h	2585	0644	edit dl rm
Stream.h	524	0644	edit dl rm
THCP.h	421	0644	edit dl rm
undef_macros.h	932	0644	edit dl rm
utils.h	871	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/cuda/comm.h (1525B)

#pragma once

#include <ATen/ATen.h>
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <ATen/cuda/ATenCUDAGeneral.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/Optional.h>

#include <cstddef>
#include <vector>

namespace torch { namespace cuda {

using tensor_list2d = std::vector<std::vector<at::Tensor>>;

TORCH_CUDA_CU_API std::vector<at::Tensor>& broadcast_out(
    const at::Tensor& tensor,
    std::vector<at::Tensor>& out_tensors);
TORCH_CUDA_CU_API std::vector<at::Tensor> broadcast(
    const at::Tensor& tensor,
    at::IntArrayRef devices);
TORCH_CUDA_CU_API tensor_list2d broadcast_coalesced(
    at::TensorList tensors,
    at::IntArrayRef devices,
    size_t buffer_size);

TORCH_CUDA_CU_API std::vector<at::Tensor>& scatter_out(
    const at::Tensor& tensor,
    std::vector<at::Tensor>& out_tensors,
    int64_t dim = 0,
    const c10::optional<std::vector<c10::optional<at::cuda::CUDAStream>>>&
        streams = c10::nullopt);

TORCH_CUDA_CU_API std::vector<at::Tensor> scatter(
    const at::Tensor& tensor,
    at::IntArrayRef devices,
    const c10::optional<std::vector<int64_t>>& chunk_sizes = c10::nullopt,
    int64_t dim = 0,
    const c10::optional<std::vector<c10::optional<at::cuda::CUDAStream>>>&
        streams = c10::nullopt);

TORCH_CUDA_CU_API at::Tensor& gather_out(
    at::TensorList tensors,
    at::Tensor& out_tensor,
    int64_t dim);

TORCH_CUDA_CU_API at::Tensor gather(
    at::TensorList tensors,
    int64_t dim,
    c10::optional<int32_t> destination_index);
}}

#pragma once #include #include #include #include #include #include #include namespace torch { namespace cuda { using tensor_list2d = std::vector>; TORCH_CUDA_CU_API std::vector& broadcast_out( const at::Tensor& tensor, std::vector& out_tensors); TORCH_CUDA_CU_API std::vector broadcast( const at::Tensor& tensor, at::IntArrayRef devices); TORCH_CUDA_CU_API tensor_list2d broadcast_coalesced( at::TensorList tensors, at::IntArrayRef devices, size_t buffer_size); TORCH_CUDA_CU_API std::vector& scatter_out( const at::Tensor& tensor, std::vector& out_tensors, int64_t dim = 0, const c10::optional>>& streams = c10::nullopt); TORCH_CUDA_CU_API std::vector scatter( const at::Tensor& tensor, at::IntArrayRef devices, const c10::optional>& chunk_sizes = c10::nullopt, int64_t dim = 0, const c10::optional>>& streams = c10::nullopt); TORCH_CUDA_CU_API at::Tensor& gather_out( at::TensorList tensors, at::Tensor& out_tensor, int64_t dim); TORCH_CUDA_CU_API at::Tensor gather( at::TensorList tensors, int64_t dim, c10::optional destination_index); }}