#ifndef GRAPH_BACKEND_DNNL_DNNL_CONSTANT_TENSOR_CACHE_HPP
#define GRAPH_BACKEND_DNNL_DNNL_CONSTANT_TENSOR_CACHE_HPP
#include "graph/interface/constant_tensor_cache.hpp"
#include "graph/backend/dnnl/common.hpp"
#include "graph/backend/dnnl/dnnl_backend.hpp"
#include "oneapi/dnnl/dnnl.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {
struct dnnl_constant_buffer_t : public graph::constant_buffer_t {
dnnl_constant_buffer_t(
size_t size, dnnl::engine &engine, graph::allocator_t *alc)
: graph::constant_buffer_t(
size, engine.get(), alc, malloc_func, free_func) {}
static void *malloc_func(
size_t size, impl::engine_t *eng, graph::allocator_t *alc) {
dnnl::engine engine;
engine.reset(eng, true); return dnnl_allocator_t::malloc(
size, engine, alc, allocator_t::mem_type_t::persistent);
}
static void free_func(
void *data, impl::engine_t *eng, graph::allocator_t *alc) {
dnnl::engine engine;
engine.reset(eng, true); if (eng->kind() == engine_kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
dnnl_allocator_t::free(data, engine, alc, ::sycl::event());
#else
dnnl_allocator_t::free(data, engine, alc);
#endif
} else if (eng->kind() == engine_kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
dnnl_allocator_t::free(data, engine, alc, ::sycl::event());
#elif DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
dnnl_allocator_t::free(data, engine, alc, cl_event());
#endif
}
}
};
inline graph::constant_tensor_cache_t::value_t dnnl_constant_cache_get_or_add(
const dnnl::engine &eng, graph::constant_tensor_cache_t::key_t key,
size_t size, const graph::constant_tensor_cache_t::value_t &value) {
auto cache = graph::get_constant_tensor_cache(
eng.get()->kind(), eng.get()->index());
assertm(cache,
"no available constant cache for specified engine kind and index");
return cache->get_or_add(
dnnl_backend_t::get_singleton().get_id(), key, size, value);
}
inline void dnnl_constant_cache_remove_if_exist(
const dnnl::engine &eng, graph::constant_tensor_cache_t::key_t key) {
auto cache = graph::get_constant_tensor_cache(
eng.get()->kind(), eng.get()->index());
assertm(cache,
"no available constant cache for specified engine kind and index");
cache->remove_if_exist(dnnl_backend_t::get_singleton().get_id(), key);
}
inline bool is_constant_cache_enabled(const dnnl::engine &eng) {
auto cache = graph::get_constant_tensor_cache(
eng.get()->kind(), eng.get()->index());
return cache && cache->get_capacity() != 0;
}
inline void dnnl_constant_cache_retain(const dnnl::engine &eng) {
auto cache = graph::get_constant_tensor_cache(
eng.get()->kind(), eng.get()->index());
cache->retain();
}
inline void dnnl_constant_cache_release(const dnnl::engine &eng) {
auto cache = graph::get_constant_tensor_cache(
eng.get()->kind(), eng.get()->index());
cache->release();
}
} } } }
#endif