#include "oneapi/dnnl/dnnl_graph.h"
#include "common/engine.hpp"
#include "graph/interface/allocator.hpp"
#include "graph/interface/tensor.hpp"
#include "graph/utils/utils.hpp"
#if DNNL_CPU_RUNTIME != DNNL_RUNTIME_SYCL
static const size_t DNNL_CPU_MEMALIGNMENT = 64;
#endif
#ifdef DNNL_WITH_SYCL
#include "oneapi/dnnl/dnnl_sycl.hpp"
static const size_t DNNL_SYCL_MEMALIGNMENT = 64;
#endif
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
#include "xpu/ocl/engine_factory.hpp"
static const size_t DNNL_OCL_MEMALIGNMENT = 0;
#endif
using namespace dnnl::impl::graph;
static void *tensor_malloc(
size_t size, const engine_t *eng, allocator_t::mem_type_t type) {
const auto *alc = static_cast<dnnl::impl::graph::allocator_t *>(
eng->get_allocator());
#ifdef DNNL_WITH_SYCL
void *dev_ptr {nullptr};
dnnl_sycl_interop_engine_get_device(const_cast<engine_t *>(eng), &dev_ptr);
auto *dev = static_cast<sycl::device *>(dev_ptr);
void *ctx_ptr {nullptr};
dnnl_sycl_interop_engine_get_context(const_cast<engine_t *>(eng), &ctx_ptr);
auto *ctx = static_cast<sycl::context *>(ctx_ptr);
#endif
if (eng->kind() == engine_kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
return alc->allocate(size, *dev, *ctx, {type, DNNL_SYCL_MEMALIGNMENT});
#else
return alc->allocate(size, {type, DNNL_CPU_MEMALIGNMENT});
#endif
} else if (eng->kind() == engine_kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
return alc->allocate(size, *dev, *ctx, {type, DNNL_SYCL_MEMALIGNMENT});
#elif DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
auto *ocl_engine = utils::downcast<
const dnnl::impl::gpu::intel::ocl::engine_t *>(eng);
const cl_device_id &ocl_dev = ocl_engine->device();
const cl_context &ocl_ctx = ocl_engine->context();
return alc->allocate(
size, ocl_dev, ocl_ctx, {type, DNNL_OCL_MEMALIGNMENT});
#else
assertm(false, "Unsupported gpu runtime");
return nullptr;
#endif
} else {
assertm(false, "Unsupported engine kind");
return nullptr;
}
}
static void tensor_free(void *p, const engine_t *eng) {
const auto *alc = static_cast<dnnl::impl::graph::allocator_t *>(
eng->get_allocator());
#ifdef DNNL_WITH_SYCL
void *dev_ptr {nullptr};
dnnl_sycl_interop_engine_get_device(const_cast<engine_t *>(eng), &dev_ptr);
auto *dev = static_cast<sycl::device *>(dev_ptr);
void *ctx_ptr {nullptr};
dnnl_sycl_interop_engine_get_context(const_cast<engine_t *>(eng), &ctx_ptr);
auto *ctx = static_cast<sycl::context *>(ctx_ptr);
#endif
if (eng->kind() == engine_kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
alc->deallocate(p, *dev, *ctx, {});
#else
alc->deallocate(p);
#endif
} else if (eng->kind() == engine_kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
alc->deallocate(p, *dev, *ctx, {});
#elif DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
auto *ocl_engine = utils::downcast<
const dnnl::impl::gpu::intel::ocl::engine_t *>(eng);
const cl_device_id &ocl_dev = ocl_engine->device();
const cl_context &ocl_ctx = ocl_engine->context();
return alc->deallocate(p, ocl_dev, ocl_ctx, {});
#else
assertm(false, "Unsupported gpu runtime");
#endif
} else {
assertm(false, "Unsupported engine kind");
}
}
dnnl_graph_tensor::dnnl_graph_tensor(
const logical_tensor_t <, const engine_t *eng, void *handle)
: lt_(lt), eng_(eng) {
if (handle == DNNL_MEMORY_ALLOCATE) {
size_t num_bytes = logical_tensor_wrapper_t(lt).size();
void *data
= tensor_malloc(num_bytes, eng, allocator_t::mem_type_t::temp);
assertm(data, "Can't allocate memory for a tensor!");
handle_.reset(data, [eng](void *p) { tensor_free(p, eng); });
} else if (lt.property == property_type::host_scalar) {
if (lt.data_type == data_type::s32) {
scalar_.s32_value = *static_cast<int32_t *>(handle);
handle_.reset(&scalar_.s32_value, dummy_destructor);
} else if (lt.data_type == data_type::f32) {
scalar_.f32_value = *static_cast<float *>(handle);
handle_.reset(&scalar_.f32_value, dummy_destructor);
} else if (lt.data_type == data_type::s64) {
scalar_.s64_value = *static_cast<int64_t *>(handle);
handle_.reset(&scalar_.s64_value, dummy_destructor);
} else {
assertm(false, "Unsupported data type for host scalar");
}
} else {
handle_.reset(handle, dummy_destructor);
}
}
status_t DNNL_API dnnl_graph_tensor_create(tensor_t **tensor,
const logical_tensor_t *logical_tensor, engine_t *eng, void *handle) {
if (utils::any_null(tensor, logical_tensor, eng))
return status::invalid_arguments;
const auto ltw = logical_tensor_wrapper_t(logical_tensor);
if (ltw.is_host_scalar()) return status::invalid_arguments;
*tensor = new tensor_t {*logical_tensor, eng, handle};
if (*tensor == nullptr) return status::out_of_memory;
if (handle == DNNL_MEMORY_ALLOCATE
&& (*tensor)->get_data_handle() == nullptr) {
delete *tensor;
*tensor = nullptr;
return status::out_of_memory;
}
return status::success;
}
status_t DNNL_API dnnl_graph_tensor_create_scalar(tensor_t **tensor,
const logical_tensor_t *logical_tensor, void *handle) {
if (utils::any_null(tensor, logical_tensor))
return status::invalid_arguments;
const auto ltw = logical_tensor_wrapper_t(logical_tensor);
if (!ltw.is_host_scalar()) return status::invalid_arguments;
if (nullptr == handle || DNNL_MEMORY_ALLOCATE == handle) {
return status::invalid_arguments;
}
*tensor = new tensor_t {*logical_tensor, nullptr, handle};
if (*tensor == nullptr) return status::out_of_memory;
return status::success;
}
status_t DNNL_API dnnl_graph_tensor_destroy(tensor_t *tensor) {
delete tensor;
return status::success;
}
status_t DNNL_API dnnl_graph_tensor_get_data_handle(
const tensor_t *tensor, void **handle) {
if (utils::any_null(tensor, handle)) return status::invalid_arguments;
const auto ltw = logical_tensor_wrapper_t(tensor->get_logical_tensor());
if (ltw.is_host_scalar()) {
*handle = nullptr;
} else {
*handle = tensor->get_data_handle();
}
return status::success;
}
status_t DNNL_API dnnl_graph_tensor_set_data_handle(
tensor_t *tensor, void *handle) {
if (tensor == nullptr) return status::invalid_arguments;
const auto ltw = logical_tensor_wrapper_t(tensor->get_logical_tensor());
if (ltw.is_host_scalar()) return status::invalid_arguments;
auto ret = tensor->set_data_handle(handle);
return ret;
}
status_t DNNL_API dnnl_graph_tensor_get_engine(
const tensor_t *tensor, engine_t **engine) {
if (utils::any_null(tensor, engine)) return status::invalid_arguments;
const auto ltw = logical_tensor_wrapper_t(tensor->get_logical_tensor());
if (ltw.is_host_scalar()) {
*engine = nullptr;
} else {
*engine = const_cast<engine_t *>(tensor->get_engine());
}
return status::success;
}
dnnl_status_t DNNL_API dnnl_graph_tensor_get_logical_tensor(
const tensor_t *tensor, logical_tensor_t *logical_tensor) {
if (utils::any_null(tensor, logical_tensor))
return status::invalid_arguments;
*logical_tensor = tensor->get_logical_tensor();
return status::success;
}