BRL-CAD/mkl-dnn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cpu: matmul: enable bias + sum support in f32 & bf16 gemm impl

Browse Source 
In bf16->f32 and f32->f32 gemm-based matmul with bias and sum
post-ops call gemm with accumulation buffer and compute sum in
pos-ops kernel to avoid falling back to reference implementation.
This commit is contained in:
Andrey Kharitonchik
2021-05-25 13:15:53 -07:00
committed by tprimak
parent 36dba73d0f
commit f993b25dbe
2 changed files with 69 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cpu/matmul/gemm_bf16_matmul.cpp
+20 -23
View File
@@ -61,23 +61,19 @@ status_t gemm_bf16_matmul_t<dst_type>::pd_t::init(engine_t *engine) {
&& gemm_based::check_gemm_compatible_formats(*this);
if (!ok) return status::unimplemented;
// set state
params_.dst_is_acc_ = dst_type == data_type::f32;
status_t status = check_and_configure_attributes();
if (status != status::success) return status;
CHECK(check_and_configure_attributes());
gemm_based::book_acc_scratchpad(*this, params_, sizeof(acc_data_t));
return status::success;
}
static bool should_gemm_execute_sum_po(
const gemm_based::params_t &params) noexcept {
static bool should_gemm_execute_sum_po(const gemm_based::params_t &params,
impl::data_type_t dst_type) noexcept {
const auto &po = params.pp_attr_.post_ops_;
static constexpr int sum_idx = 0;
return po.len() > 0 && po.contain(primitive_kind::sum, sum_idx)
&& params.dst_is_acc_;
&& dst_type == data_type::f32 && params.gemm_applies_output_scales_;
}
template <impl::data_type_t dst_type>
@@ -91,11 +87,7 @@ status_t gemm_bf16_matmul_t<dst_type>::pd_t::check_and_configure_attributes() {
auto check_attr_post_ops = [&]() -> bool {
using namespace primitive_kind;
const auto &post_ops = attr()->post_ops_;
if (IMPLICATION(post_ops.contain(sum, 0),
params_.gemm_applies_output_scales_)) {
return cpu::inner_product_utils::post_ops_ok(post_ops, dst_md());
}
return false;
return cpu::inner_product_utils::post_ops_ok(post_ops, dst_md());
};
// check basic attributes
@@ -105,19 +97,24 @@ status_t gemm_bf16_matmul_t<dst_type>::pd_t::check_and_configure_attributes() {
CHECK(params_.pp_attr_.copy_from(*attr()));
params_.gemm_applies_output_scales_
= attr()->output_scales_.mask_ == 0 && !with_bias();
if (params_.gemm_applies_output_scales_)
params_.pp_attr_.output_scales_.set(1.f);
// check post-ops
if (check_attr_post_ops()) {
if (should_gemm_execute_sum_po(params_)) {
// set state
const auto &po = params_.pp_attr_.post_ops_;
static constexpr int sum_idx = 0;
params_.gemm_beta_ = po.entry_[sum_idx].sum.scale;
}
} else {
return status::unimplemented;
if (!check_attr_post_ops()) return status::unimplemented;
const bool sum_po_via_gemm_beta
= should_gemm_execute_sum_po(params_, dst_type);
// set state
params_.dst_is_acc_ = dst_type == data_type::f32
&& IMPLICATION(attr()->post_ops_.find(primitive_kind::sum) != -1,
sum_po_via_gemm_beta);
if (sum_po_via_gemm_beta) {
// set state
const auto &po = params_.pp_attr_.post_ops_;
static constexpr int sum_idx = 0;
params_.gemm_beta_ = po.entry_[sum_idx].sum.scale;
}
// set state
@@ -129,7 +126,7 @@ status_t gemm_bf16_matmul_t<dst_type>::pd_t::check_and_configure_attributes() {
template <impl::data_type_t dst_type>
bool gemm_bf16_matmul_t<dst_type>::should_skip_sum_po() const noexcept {
return should_gemm_execute_sum_po(pd()->params());
return should_gemm_execute_sum_po(pd()->params(), dst_type);
}
template <impl::data_type_t dst_type>
src/cpu/matmul/gemm_f32_matmul.cpp
+49 -19
View File
@@ -58,14 +58,16 @@ status_t gemm_f32_matmul_t::pd_t::init(engine_t *engine) {
if (!ok) return status::unimplemented;
// set state
params_.dst_is_acc_ = true;
if (!has_runtime_dims_or_strides())
params_.can_fuse_src_batch_dims_
= matmul_helper_t(src_md(), weights_md(), dst_md())
.can_fuse_src_batch_dims();
return check_and_configure_attributes();
CHECK(check_and_configure_attributes());
gemm_based::book_acc_scratchpad(*this, params_, sizeof(acc_data_t));
return status::success;
}
static bool should_gemm_execute_sum_po(
@@ -73,7 +75,7 @@ static bool should_gemm_execute_sum_po(
const auto &po = params.pp_attr_.post_ops_;
static constexpr int sum_idx = 0;
return po.len() > 0 && po.contain(primitive_kind::sum, sum_idx)
&& params.dst_is_acc_;
&& params.gemm_applies_output_scales_;
}
status_t gemm_f32_matmul_t::pd_t::check_and_configure_attributes() {
@@ -86,11 +88,7 @@ status_t gemm_f32_matmul_t::pd_t::check_and_configure_attributes() {
auto check_attr_post_ops = [&]() -> bool {
using namespace primitive_kind;
const auto &post_ops = attr()->post_ops_;
if (IMPLICATION(post_ops.contain(sum, 0),
params_.gemm_applies_output_scales_)) {
return cpu::inner_product_utils::post_ops_ok(post_ops, dst_md());
}
return false;
return cpu::inner_product_utils::post_ops_ok(post_ops, dst_md());
};
// check basic attributes
@@ -105,8 +103,13 @@ status_t gemm_f32_matmul_t::pd_t::check_and_configure_attributes() {
// check post-ops
if (!check_attr_post_ops()) return status::unimplemented;
const bool sum_po_via_gemm_beta = should_gemm_execute_sum_po(params_);
// set state
params_.dst_is_acc_
= IMPLICATION(attr()->post_ops_.find(primitive_kind::sum) != -1,
sum_po_via_gemm_beta);
if (should_gemm_execute_sum_po(params_)) {
if (sum_po_via_gemm_beta) {
// set state
const auto &po = params_.pp_attr_.post_ops_;
static constexpr int sum_idx = 0;
@@ -114,8 +117,8 @@ status_t gemm_f32_matmul_t::pd_t::check_and_configure_attributes() {
}
// set state
params_.has_pp_kernel_
= with_bias() || !params_.pp_attr_.has_default_values();
params_.has_pp_kernel_ = !params_.dst_is_acc_ || with_bias()
|| !params_.pp_attr_.has_default_values();
return status::success;
}
@@ -159,6 +162,27 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
? helper.can_fuse_src_batch_dims()
: params.can_fuse_src_batch_dims_;
const dim_t acc_stride = gemm_based::get_scratchpad_size(
batch, M, N, can_fuse_src_batch_dims);
bool dst_is_acc = params.dst_is_acc_;
acc_data_t *acc = dst_is_acc
? (acc_data_t *)dst
: ctx.get_scratchpad_grantor().template get<acc_data_t>(
memory_tracking::names::key_matmul_dst_in_acc_dt);
// case: dynamic sizes
bool need_free_acc = false;
if (acc == nullptr) {
acc = (acc_data_t *)malloc(sizeof(acc_data_t) * acc_stride
* ((can_fuse_src_batch_dims || batch == 1)
? 1
: (dim_t)dnnl_get_max_threads()),
64);
if (acc == nullptr) return status::out_of_memory;
need_free_acc = true;
}
const dim_t acc_ldc = dst_is_acc ? ldc : N;
std::atomic<status_t> st(status::success);
const bool parallel_over_batch = batch > 1 && !can_fuse_src_batch_dims;
if (parallel_over_batch) {
@@ -178,6 +202,9 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
dim_t cur_b {0}, cur_m {0}, cur_n {0};
dims_t s_dims_idx, w_dims_idx, d_dims_idx;
size_t i_work = t_work_start;
const bool reuse_acc = acc != (acc_data_t *)dst;
acc_data_t *curr_acc
= reuse_acc ? acc + ithr * acc_stride : nullptr;
while (i_work < t_work_end) {
utils::nd_iterator_init(
@@ -199,7 +226,9 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
const src_data_t *curr_src = src + src_d.off_v(s_dims_idx);
const weights_data_t *curr_weights
= weights + weights_d.off_v(w_dims_idx);
dst_data_t *curr_dst = dst + dst_d.off_v(d_dims_idx);
const dim_t dst_off = dst_d.off_v(d_dims_idx);
dst_data_t *curr_dst = dst + dst_off;
if (!reuse_acc) curr_acc = acc + dst_off;
dim_t gemm_M {0}, gemm_N {0};
const size_t rem_work = t_work_end - i_work;
@@ -220,7 +249,7 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
status_t st_thr = extended_sgemm(&transB, &transA, &gemm_N,
&gemm_M, &K, &alpha, curr_weights, &ldb, curr_src, &lda,
&beta, curr_dst, &ldc, nullptr, false);
&beta, curr_acc, &acc_ldc, nullptr, false);
if (st_thr != status::success) {
st = st_thr;
return;
@@ -229,7 +258,7 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
if (params.has_pp_kernel_) {
const float *pp_scales
= params.get_post_processing_scales(scales);
(*pp_kernel_)(curr_dst, curr_dst,
(*pp_kernel_)(curr_dst, curr_acc,
bias
+ static_cast<ptrdiff_t>(i_work % N)
* bia_dt_size,
@@ -246,22 +275,23 @@ status_t gemm_f32_matmul_t::execute_ref(const exec_ctx_t &ctx) const {
M = batch * M;
st = extended_sgemm(&transB, &transA, &N, &M, &K, &alpha, weights, &ldb,
src, &lda, &beta, dst, &ldc, nullptr, false);
if (st != status::success) return st;
src, &lda, &beta, acc, &acc_ldc, nullptr, false);
if (params.has_pp_kernel_) {
if (st == status::success && params.has_pp_kernel_) {
const bool force_sequential = pp_kernel_->sequential_kernel();
const float *pp_scales = params.get_post_processing_scales(scales);
parallel(force_sequential ? 1 : 0, [&](int ithr, int nthr) {
size_t start {}, end {};
balance211((size_t)(M * N), nthr, ithr, start, end);
(*pp_kernel_)(dst, dst, bias, pp_scales, start, end, (size_t)N,
(*pp_kernel_)(dst, acc, bias, pp_scales, start, end, (size_t)N,
ldc, nullptr, post_ops_binary_rhs_arg_vec.data(), dst,
ctx, *pd()->dst_md());
});
}
}
if (need_free_acc) free(acc);
return st;
}