memory.h

Go to the documentation of this file.

#ifndef ATEN_KERNELS_MEMORY_H_
#define ATEN_KERNELS_MEMORY_H_
 
#include <vector>
#include <complex>
 
#include <base/macros/macros.h>
#include <ATen/core/tensor_types.h>
 
namespace container {
namespace kernels {
 
template <typename T, typename Device>
struct resize_memory {
    void operator()(T*& arr, const size_t& size, const char* record_in = nullptr);
};
 
template <typename T, typename Device>
struct set_memory {
    void operator()(T* arr, const T& var, const size_t& size);
};
 
template <typename T, typename Device_out, typename Device_in>
struct synchronize_memory {
    void operator()(
        T* arr_out,
        const T* arr_in,
        const size_t& size);
};
 
template <typename T, typename Device_out, typename Device_in>
struct synchronize_memory_stride {
    void operator()(
        T* arr_out,
        const T* arr_in,
        const std::vector<int64_t>& out_size,
        const std::vector<int64_t>& in_size)
    {
        REQUIRES_OK(in_size.size() == out_size.size() && in_size.size() <= 2);
        if (in_size.size() == 1) {
            synchronize_memory<T, Device_out, Device_in>()(arr_out, arr_in, in_size[0]);
        }
        else {
            for (int64_t ii = 0; ii < out_size[0]; ii++) {
                synchronize_memory<T, Device_out, Device_in>()(arr_out + ii * out_size[1], arr_in + ii * in_size[1], in_size[1]);
            }
        }
    }
};
 
template <typename T_out, typename T_in, typename Device_out, typename Device_in>
struct cast_memory {
    void operator()(
        T_out* arr_out,
        const T_in* arr_in,
        const size_t& size);
};
 
 
template <typename T, typename Device>
struct delete_memory {
    void operator()(T* arr);
};
 
} // namespace kernels
} // namespace container
 
#endif // ATEN_KERNELS_MEMORY_H_