mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-22 16:17:37 +00:00
72 lines
2.3 KiB
C++
72 lines
2.3 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#pragma once
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/host/host_tensor.hpp"
|
|
#include <thread>
|
|
|
|
namespace ck_tile {
|
|
|
|
template <typename ADataType,
|
|
typename CompDataType,
|
|
typename BDataType,
|
|
typename CompElementOp = ck_tile::identity>
|
|
CK_TILE_HOST void reference_batched_softmax(
|
|
const HostTensor<ADataType>& a_b_m_n,
|
|
HostTensor<BDataType>& b_b_m_n,
|
|
const CompElementOp& comp_element_op = {},
|
|
std::optional<std::reference_wrapper<HostTensor<CompDataType>>> lse_b_m = std::nullopt)
|
|
{
|
|
const int N = a_b_m_n.mDesc.get_lengths()[2];
|
|
|
|
auto f = [&](auto batch, auto m) {
|
|
CompDataType v_max = -ck_tile::numeric<CompDataType>::infinity();
|
|
|
|
// max
|
|
for(int n = 0; n < N; ++n)
|
|
{
|
|
const CompDataType v_a = ck_tile::type_convert<CompDataType>(a_b_m_n(batch, m, n));
|
|
|
|
v_max = v_max < v_a ? v_a : v_max;
|
|
}
|
|
|
|
CompDataType v_exp_sum = 0;
|
|
// validate v_max if all the elements within a row are -INF
|
|
if(std::isinf(v_max) && v_max < 0)
|
|
{
|
|
v_max = ck_tile::type_convert<CompDataType>(0.f);
|
|
}
|
|
|
|
// sum
|
|
for(int n = 0; n < N; ++n)
|
|
{
|
|
const CompDataType v_a = ck_tile::type_convert<CompDataType>(a_b_m_n(batch, m, n));
|
|
|
|
v_exp_sum += ck_tile::exp(v_a - v_max);
|
|
}
|
|
|
|
// if sum is zero(masked), or nan/inf(other computation error), don't do divide
|
|
CompDataType inv_sum = (v_exp_sum == 0.f ? 1.f : 1.f / v_exp_sum);
|
|
|
|
// elementwise
|
|
for(int n = 0; n < N; ++n)
|
|
{
|
|
const CompDataType v_a = ck_tile::type_convert<CompDataType>(a_b_m_n(batch, m, n));
|
|
const CompDataType v_b = ck_tile::exp(v_a - v_max) * inv_sum;
|
|
|
|
b_b_m_n(batch, m, n) = ck_tile::type_convert<BDataType>(comp_element_op(v_b));
|
|
}
|
|
// lse
|
|
if(lse_b_m)
|
|
{
|
|
lse_b_m->get()(batch, m) = v_max + ck_tile::log(v_exp_sum);
|
|
}
|
|
};
|
|
|
|
make_ParallelTensorFunctor(f, b_b_m_n.mDesc.get_lengths()[0], b_b_m_n.mDesc.get_lengths()[1])(
|
|
std::thread::hardware_concurrency());
|
|
}
|
|
} // namespace ck_tile
|