fix build at ubuntu 22.04

MegEngine · Mar 20, 2023 · 0d659f3 · 0d659f3
1 parent 4ade75c
commit 0d659f3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/src/cpu/rnn/jit_uni_rnn_common_postgemm_dispatcher.hpp b/src/cpu/rnn/jit_uni_rnn_common_postgemm_dispatcher.hpp
@@ -232,7 +232,7 @@ using rnn_postgemm_bwd_f32_t = rnn_postgemm_dispatcher<prop_kind::backward,
 using rnn_postgemm_fwd_bf16_t = rnn_postgemm_dispatcher<prop_kind::forward,
         data_type::bf16, data_type::f32>;
 using rnn_postgemm_bwd_bf16_t = rnn_postgemm_dispatcher<prop_kind::backward,
-        data_type::bf16, data_type::f32>;
+        data_type::bf16, data_type::bf16>;
 using rnn_postgemm_fwd_u8_t = rnn_postgemm_dispatcher<prop_kind::forward,
         data_type::u8, data_type::s32>;
 

diff --git a/src/cpu/rnn/ref_rnn.hpp b/src/cpu/rnn/ref_rnn.hpp
@@ -61,12 +61,13 @@ void gates_reduction(const rnn_utils::rnn_conf_t &rnn,
 template <prop_kind_t aprop, impl::data_type_t src_type,
         impl::data_type_t weights_type, impl::data_type_t acc_type>
 struct _ref_rnn_common_t : public primitive_impl_t {
+    static constexpr impl::data_type_t scratch_type
+            = aprop == prop_kind::forward ? acc_type : src_type;
+
     typedef typename prec_traits<src_type>::type src_data_t;
     typedef typename prec_traits<weights_type>::type weights_data_t;
     typedef typename prec_traits<acc_type>::type acc_data_t;
-
-    typedef typename utils::conditional<aprop == prop_kind::forward, acc_data_t,
-            src_data_t>::type scratch_data_t;
+    typedef typename prec_traits<scratch_type>::type scratch_data_t;
 
     using class_name
             = _ref_rnn_common_t<aprop, src_type, weights_type, acc_type>;
@@ -224,8 +225,9 @@ struct _ref_rnn_common_t : public primitive_impl_t {
         set_gemm_funcs(pd()->rnn_.use_layer_packed_gemm, gemm_layer_func,
                 weights_layer_assign_func);
 
-        rnn_postgemm_ = new rnn_postgemm_dispatcher<aprop, src_type, acc_type>(
-                pd()->rnn_, pd());
+        rnn_postgemm_
+                = new rnn_postgemm_dispatcher<aprop, src_type, scratch_type>(
+                        pd()->rnn_, pd());
         assert(rnn_postgemm_ != nullptr);
         switch (pd()->cell_kind()) {
             case alg_kind::vanilla_rnn:
@@ -308,7 +310,7 @@ struct _ref_rnn_common_t : public primitive_impl_t {
     size_t ws_grid_comp_offset_;
     size_t scratch_gates_offset_;
     size_t scratch_cell_offset_;
-    rnn_postgemm_dispatcher<aprop, src_type, acc_type> *rnn_postgemm_;
+    rnn_postgemm_dispatcher<aprop, src_type, scratch_type> *rnn_postgemm_;
 
     grid_execution_f grid_computation;
     cell_execution_f cell_func;