#ifndef CPU_CPU_STREAM_HPP
#define CPU_CPU_STREAM_HPP
#include "oneapi/dnnl/dnnl_config.h"
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
#include "oneapi/dnnl/dnnl_threadpool_iface.hpp"
#endif
#include "common/c_types_map.hpp"
#include "common/dnnl_thread.hpp"
#include "common/stream.hpp"
namespace dnnl {
namespace impl {
namespace cpu {
struct cpu_stream_t : public stream_t {
cpu_stream_t(engine_t *engine, impl::stream_impl_t *stream_impl)
: stream_t(engine, stream_impl) {}
~cpu_stream_t() override = default;
dnnl::impl::status_t wait() override {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
dnnl::threadpool_interop::threadpool_iface *tp;
auto rc = this->get_threadpool(&tp);
if (rc == status::success && tp) {
if (tp->get_flags()
& threadpool_interop::threadpool_iface::ASYNCHRONOUS)
tp->wait();
}
#endif
return dnnl::impl::status::success;
}
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
cpu_stream_t(engine_t *engine,
dnnl::threadpool_interop::threadpool_iface *threadpool)
: stream_t(engine, new impl::stream_impl_t(threadpool)) {}
void before_exec_hook() override {
dnnl::threadpool_interop::threadpool_iface *tp;
auto rc = this->get_threadpool(&tp);
if (rc == status::success) threadpool_utils::activate_threadpool(tp);
}
void after_exec_hook() override {
threadpool_utils::deactivate_threadpool();
}
#endif
};
} } }
#endif