diff --git a/include/armadillo_bits/op_sum_bones.hpp b/include/armadillo_bits/op_sum_bones.hpp index f687c9b0..2445990a 100644 --- a/include/armadillo_bits/op_sum_bones.hpp +++ b/include/armadillo_bits/op_sum_bones.hpp @@ -40,6 +40,9 @@ struct op_sum template inline static void apply_mat_noalias(Mat& out, const Mat& X, const uword dim); + template + inline static void apply_mat_noalias_promote(Mat& out, const Mat& X, const uword dim); + template inline static void apply_mat_square_noalias(Mat& out, const Mat& X, const uword dim); diff --git a/include/armadillo_bits/op_sum_meat.hpp b/include/armadillo_bits/op_sum_meat.hpp index c60fd921..876d47f8 100644 --- a/include/armadillo_bits/op_sum_meat.hpp +++ b/include/armadillo_bits/op_sum_meat.hpp @@ -69,7 +69,7 @@ op_sum::apply(Mat& out, const Op< eOp, op return; } - + op_sum::apply_generic(out, in); } @@ -120,7 +120,7 @@ op_sum::apply_generic(Mat& out, const Op& in) arma_conform_check( (dim > 1), "sum(): parameter 'dim' must be 0 or 1" ); - if((is_Mat::value) || (is_Mat::stored_type>::value) || (arma_config::openmp && Proxy::use_mp)) + if((is_Mat::value) || (is_Mat::stored_type>::value) || (arma_config::openmp && Proxy::use_mp) || (is_fp16::yes) || (is_cx_fp16::yes)) { const quasi_unwrap U(in.m); @@ -165,6 +165,12 @@ op_sum::apply_mat_noalias(Mat& out, const Mat& X, const uword dim) { arma_debug_sigprint(); + #if defined(ARMA_HAVE_FP16) + { + if(is_fp16::yes || is_cx_fp16::yes) { op_sum::apply_mat_noalias_promote(out, X, dim); return; } + } + #endif + const uword X_n_rows = X.n_rows; const uword X_n_cols = X.n_cols; @@ -204,6 +210,47 @@ op_sum::apply_mat_noalias(Mat& out, const Mat& X, const uword dim) +template +inline +void +op_sum::apply_mat_noalias_promote(Mat& out, const Mat& X, const uword dim) + { + arma_debug_sigprint(); + + const uword X_n_rows = X.n_rows; + const uword X_n_cols = X.n_cols; + + const uword out_n_rows = (dim == 0) ? uword(1) : X_n_rows; + const uword out_n_cols = (dim == 0) ? X_n_cols : uword(1); + + out.set_size(out_n_rows, out_n_cols); + + if(X.n_elem == 0) { out.zeros(); return; } + + eT* out_mem = out.memptr(); + + if(dim == 0) + { + for(uword col=0; col < X_n_cols; ++col) + { + out_mem[col] = arrayops::accumulate_promote( X.colptr(col), X_n_rows ); + } + } + else + { + podarray tmp; + + for(uword row=0; row < X_n_rows; ++row) + { + tmp.copy_row(X, row); + + out_mem[row] = arrayops::accumulate_promote( tmp.memptr(), tmp.n_elem ); + } + } + } + + + template inline void