14 template <
typename OutDataType,
typename AccDataType>
18 for(
int n = 0; n < N; ++n)
20 o(m, n) = ck_tile::type_convert<OutDataType>(acc(m, n));
24 template <
typename OutDataType,
typename AccDataType>
33 template <
typename XDataType,
34 typename GammaDataType,
35 typename ComputeDataType,
37 typename InvRmsDataType,
38 typename UnquantYDataType,
39 typename Epilogue = reference_rmsnorm2d_default_epilogue>
45 ComputeDataType epsilon,
46 Epilogue epilogue_functor = {})
48 auto rmsnorm2d_fwd_func = [&](
auto m) {
51 ComputeDataType mean_square = 0;
52 ComputeDataType divisor = 0;
54 for(
int n = 0; n < N; ++n)
56 ComputeDataType x = ck_tile::type_convert<ComputeDataType>(x_m_n(m, n));
60 mean_square = mean_square / N;
61 divisor = ck_tile::type_convert<ComputeDataType>(1) /
ck_tile::sqrt(mean_square + epsilon);
63 if constexpr(!std::is_same_v<InvRmsDataType, ck_tile::null_type>)
64 invRms_m(m) = ck_tile::type_convert<InvRmsDataType>(divisor);
67 for(
int n = 0; n < N; ++n)
69 ComputeDataType x = ck_tile::type_convert<ComputeDataType>(x_m_n(m, n));
70 ComputeDataType gamma = ck_tile::type_convert<ComputeDataType>(gamma_n(n));
71 acc(m, n) = x * divisor * gamma;
74 if constexpr(!std::is_same_v<UnquantYDataType, ck_tile::null_type>)
76 epilogue_functor(m, unquant_y_m_n, y_m_n, acc);
80 epilogue_functor(m, y_m_n, acc);
85 std::thread::hardware_concurrency());
Definition: cluster_descriptor.hpp:13
CK_TILE_HOST auto make_ParallelTensorFunctor(F f, Xs... xs)
Definition: host_tensor.hpp:329
CK_TILE_DEVICE bfloat16_t sqrt(bfloat16_t x)
Definition: bfloat16.hpp:417
void reference_rmsnorm2d_fwd(const HostTensor< XDataType > &x_m_n, const HostTensor< GammaDataType > &gamma_n, HostTensor< YDataType > &y_m_n, HostTensor< InvRmsDataType > &invRms_m, HostTensor< UnquantYDataType > &unquant_y_m_n, ComputeDataType epsilon, Epilogue epilogue_functor={})
Definition: reference_rmsnorm2d_fwd.hpp:40
const std::vector< std::size_t > & get_lengths() const
Definition: host_tensor.hpp:198
Definition: host_tensor.hpp:336
decltype(auto) get_lengths() const
Definition: host_tensor.hpp:390
decltype(auto) get_strides() const
Definition: host_tensor.hpp:394
Descriptor mDesc
Definition: host_tensor.hpp:800
Definition: reference_rmsnorm2d_fwd.hpp:13
auto operator()(int m, const HostTensor< AccDataType > &acc)
Definition: reference_rmsnorm2d_fwd.hpp:25
void operator()(int m, HostTensor< OutDataType > &o, const HostTensor< AccDataType > &acc)
Definition: reference_rmsnorm2d_fwd.hpp:15