Repository URL to install this package:
|
Version:
2.4.0 ▾
|
// NOTE: Like interface.cpp, this file will be copied into AOTInductor
// generated output. This file is intended to keep implementation
// details separate from the implementation of the AOTI public
// interface. Note also that #includes should go into interface.cpp
// for simplicity of maintenance.
namespace torch {
namespace aot_inductor {
template <typename T>
void convert_output_to_handle(
const ArrayRefTensor<T>& output,
AtenTensorHandle& handle) {
handle = output.expensiveCopyToTensor();
}
template <typename... Ts, std::size_t... Is>
void convert_outputs_to_handles_helper(
const std::tuple<ArrayRefTensor<Ts>...>& outputs,
AtenTensorHandle* output_handles,
std::index_sequence<Is...>) {
(convert_output_to_handle(std::get<Is>(outputs), output_handles[Is]), ...);
}
template <typename... Ts>
void convert_outputs_to_handles(
const std::tuple<ArrayRefTensor<Ts>...>& outputs,
AtenTensorHandle* output_handles) {
convert_outputs_to_handles_helper(
outputs, output_handles, std::make_index_sequence<sizeof...(Ts)>());
}
template <typename T>
void convert_handle_to_arrayref_tensor(
AtenTensorHandle handle,
ArrayRefTensor<T>& input) {
void* data_ptr;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle, &data_ptr));
int64_t dim;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_dim(handle, &dim));
int64_t numel;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_numel(handle, &numel));
int64_t* sizes;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle, &sizes));
int64_t* strides;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle, &strides));
int32_t dtype;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(handle, &dtype));
int32_t device_type;
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_device_type(handle, &device_type));
int32_t device_index;
AOTI_TORCH_ERROR_CODE_CHECK(
aoti_torch_get_device_index(handle, &device_index));
input = ArrayRefTensor<T>(
MiniArrayRef<T>(reinterpret_cast<T*>(data_ptr), numel),
MiniArrayRef<const int64_t>(sizes, dim),
MiniArrayRef<const int64_t>(strides, dim),
device_type,
device_index);
}
template <typename... Ts, std::size_t... Is>
void convert_handles_to_inputs_helper(
AtenTensorHandle* input_handles,
std::tuple<ArrayRefTensor<Ts>...>& inputs,
std::index_sequence<Is...>) {
(convert_handle_to_arrayref_tensor(input_handles[Is], std::get<Is>(inputs)),
...);
}
template <typename... Ts>
void convert_handles_to_inputs(
AtenTensorHandle* input_handles,
std::tuple<ArrayRefTensor<Ts>...>& inputs) {
convert_handles_to_inputs_helper(
input_handles, inputs, std::make_index_sequence<sizeof...(Ts)>());
}
template <typename T>
void assert_numel(const ArrayRefTensor<T>& tensor, int64_t numel) {
if (tensor.numel() != numel) {
std::stringstream err;
err << "incorrect numel for input tensor. expected " << numel << ", got " << tensor.numel();
throw std::runtime_error(err.str());
}
}
} // namespace aot_inductor
} // namespace torch