#include "gpu/intel/softmax/xe.hpp"
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace softmax {
status_t xe_fwd_t::execute_generic(const exec_ctx_t &ctx) const {
if (pd()->has_zero_dim_memory()) return status::success;
auto &src = CTX_IN_STORAGE(DNNL_ARG_SRC);
auto &dst = CTX_OUT_STORAGE(DNNL_ARG_DST);
auto &src_scale = CTX_IN_STORAGE(DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC);
auto &dst_scale = CTX_IN_STORAGE(DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST);
compute::kernel_arg_list_t arg_list;
arg_list.set(0, src);
arg_list.set(1, dst);
arg_list.set(2, src_scale);
arg_list.set(3, dst_scale);
auto nd_range = compute::nd_range_t(pd()->gws, pd()->lws);
return parallel_for(ctx, nd_range, kernel_, arg_list);
}
status_t xe_bwd_t::execute_generic(const exec_ctx_t &ctx) const {
if (pd()->has_zero_dim_memory()) return status::success;
auto &dst = CTX_IN_STORAGE(DNNL_ARG_DST);
auto &diff_dst = CTX_IN_STORAGE(DNNL_ARG_DIFF_DST);
auto &diff_src = CTX_OUT_STORAGE(DNNL_ARG_DIFF_SRC);
compute::kernel_arg_list_t arg_list;
arg_list.set(0, dst);
arg_list.set(1, diff_src);
arg_list.set(2, diff_dst);
auto nd_range = compute::nd_range_t(pd()->gws, pd()->lws);
return parallel_for(ctx, nd_range, kernel_, arg_list);
}
} } } } }