mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
8f0603b128
Summary: * adds TORCH_API and AT_CUDA_API in places * refactor code generation Python logic to separate caffe2/torch outputs * fix hip and asan * remove profiler_cuda from hip * fix gcc warnings for enums * Fix PythonOp::Kind Pull Request resolved: https://github.com/pytorch/pytorch/pull/19554 Differential Revision: D15082727 Pulled By: kostmo fbshipit-source-id: 83a8a99717f025ab44b29608848928d76b3147a4
77 lines
2.4 KiB
C++
77 lines
2.4 KiB
C++
#pragma once
|
|
|
|
#include <ATen/core/functional.h>
|
|
#include <torch/csrc/WindowsTorchApiMacro.h>
|
|
#include <ATen/ATen.h>
|
|
#include <utility>
|
|
|
|
namespace torch { namespace utils {
|
|
|
|
inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
|
|
static auto flatten = [](const at::Tensor &t) { return t.contiguous().view({-1}); };
|
|
if (tensors.size() == 1)
|
|
return flatten(tensors[0]);
|
|
return at::cat(fmap(tensors, flatten));
|
|
}
|
|
|
|
inline std::vector<at::Tensor> unflatten_dense_tensors(const at::Tensor& flat, at::TensorList tensors) {
|
|
std::vector<at::Tensor> outputs;
|
|
outputs.reserve(tensors.size());
|
|
size_t offset = 0;
|
|
for (const auto & tensor : tensors) {
|
|
auto numel = tensor.numel();
|
|
outputs.push_back(flat.narrow(0, offset, numel).view(tensor.sizes()));
|
|
offset += numel;
|
|
}
|
|
return outputs;
|
|
}
|
|
|
|
|
|
struct TensorGroup {
|
|
std::vector<at::Tensor> tensors;
|
|
size_t size = 0;
|
|
|
|
at::DeprecatedTypeProperties& type() {
|
|
AT_ASSERT(!tensors.empty());
|
|
return tensors[0].type();
|
|
}
|
|
};
|
|
|
|
// Helper function that takes a list of tensors and splits them into tensor
|
|
// groups by the size limit and outputs these tensor groups. If the input
|
|
// tensors are of different tensor types, they will be split into different
|
|
// groups as well.
|
|
//
|
|
// Two options of splitting provided to the user,
|
|
//
|
|
// Imagine the size_limit is 256 and the list of input tensors are:
|
|
// tensor_a(fp16 - 128 bytes),
|
|
// tensor_b(fp32 - 256 bytes),
|
|
// tensor_c(fp16 - 128 bytes),
|
|
//
|
|
// when fine_grained == false:
|
|
// The function will read the list of tensors sequentially and accumulate
|
|
// enough tensors for each data type until the size_limit, therefore:
|
|
// it will output: {{tensor_a, tensor_c}, {tensor_b}}
|
|
//
|
|
// when fine_grained == true:
|
|
// The function will read the list of tensors sequentially and accumulate
|
|
// enough tensors for all data types until the size_limit, and then split
|
|
// the accumulated tensors into different groups by data types, therefore:
|
|
// it will output: {{tensor_a}, {tensor_b}, {tensor_c}}
|
|
TORCH_API std::vector<TensorGroup> take_tensors(
|
|
at::TensorList tensors,
|
|
size_t size_limit,
|
|
bool fine_grained = false);
|
|
|
|
TORCH_API void reorder_tensors_like(std::vector<at::Tensor>& tensors, at::TensorList order);
|
|
|
|
TORCH_API std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(at::TensorList tensors);
|
|
|
|
TORCH_API std::vector<at::Tensor> unflatten_sparse_tensors(
|
|
const at::Tensor& flat_indices,
|
|
const at::Tensor& flat_values,
|
|
at::TensorList tensors);
|
|
|
|
}}
|