#include "gpu/intel/lrn/ref.hpp"
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace lrn {
status_t ref_fwd_t::execute_forward(const exec_ctx_t &ctx) const {
auto &src = CTX_IN_STORAGE(DNNL_ARG_SRC);
auto &dst = CTX_OUT_STORAGE(DNNL_ARG_DST);
auto &ws = CTX_OUT_STORAGE(DNNL_ARG_WORKSPACE);
compute::kernel_arg_list_t arg_list;
arg_list.set(0, src);
if (pd()->desc()->prop_kind == prop_kind::forward_training) {
arg_list.set(1, ws);
arg_list.set(2, dst);
} else {
arg_list.set(1, dst);
}
auto nd_range = pd()->dispatch.nd_range();
status_t status = parallel_for(ctx, nd_range, kernel_, arg_list);
return status;
}
status_t ref_bwd_t::execute_backward(const exec_ctx_t &ctx) const {
auto &src = CTX_IN_STORAGE(DNNL_ARG_SRC);
auto &diff_dst = CTX_IN_STORAGE(DNNL_ARG_DIFF_DST);
auto &ws = CTX_IN_STORAGE(DNNL_ARG_WORKSPACE);
auto &diff_src = CTX_OUT_STORAGE(DNNL_ARG_DIFF_SRC);
compute::kernel_arg_list_t arg_list;
arg_list.set(0, src);
arg_list.set(1, diff_dst);
arg_list.set(2, ws);
arg_list.set(3, diff_src);
auto nd_range = pd()->dispatch.nd_range();
status_t status = parallel_for(ctx, nd_range, kernel_, arg_list);
return status;
}
} } } } }