mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-05 14:11:29 +00:00
0d9118226b
* Add NumReduceDim template parameter to DeviceSoftmax and Softmax client API to simplify instances collecting * Move the generic kernel instance to be the first of the instance list for elementwise op of normalization * Add GetGenericInstance() interface for DeviceOperationInstanceFactory class of DeviceSoftmax * Add testing of GetGenericInstance() in client_example of Softmax * Revert "Add testing of GetGenericInstance() in client_example of Softmax" This reverts commitf629cd9a93. * Revert "Add GetGenericInstance() interface for DeviceOperationInstanceFactory class of DeviceSoftmax" This reverts commita9f0d000eb. * Support generic kernel instance to be the first instance returned by GetInstances() for GroupNorm * Move generic kernel instance to separate tuple for elementwise op of normalization * Remove un-used files for softmax instance * Store generic kernel instance to separate tuple for softmax * Add IsSupported checking for generic instance to client example of softmax * Replace the get_device_normalize_from_mean_meansquare_instances() by the DeviceOperationInstanceFactory class for elementwise-normalization * clang-format fix * Remove int8 from softmax instances --------- Co-authored-by: zjing14 <zhangjing14@gmail.com>
227 lines
12 KiB
C++
227 lines
12 KiB
C++
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
|
|
|
#pragma once
|
|
|
|
#include <string>
|
|
#include <sstream>
|
|
#include <tuple>
|
|
#include <vector>
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "ck/ck.hpp"
|
|
#include "ck/tensor_operation/gpu/device/impl/device_softmax_impl.hpp"
|
|
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
|
|
#include "include/ck/utility/data_type.hpp"
|
|
#include "profiler/profile_softmax_impl.hpp"
|
|
|
|
namespace ck {
|
|
|
|
template <typename Range>
|
|
std::string serialize_range(const Range& range)
|
|
{
|
|
std::stringstream ss;
|
|
for(auto& r : range)
|
|
{
|
|
ss << r << ", ";
|
|
}
|
|
std::string str = ss.str();
|
|
return std::string(str.begin(), str.end() - 2);
|
|
}
|
|
|
|
template <typename Tuple>
|
|
class TestSoftmax : public ::testing::Test
|
|
{
|
|
protected:
|
|
using InDataType = std::tuple_element_t<0, Tuple>;
|
|
using AccDataType = std::tuple_element_t<1, Tuple>;
|
|
using OutDataType = std::tuple_element_t<2, Tuple>;
|
|
static constexpr index_t Rank = std::tuple_element_t<3, Tuple>{}.value;
|
|
|
|
public:
|
|
std::vector<std::vector<index_t>> in_lengths_ = {{2, 128, 1024}, {4, 16, 8448}, {128, 128, 64}};
|
|
std::vector<std::vector<AccDataType>> scales_ = {{2, 0}, {0, 2}, {2, 2}};
|
|
bool bench_ = false; // measure kernel performance
|
|
bool verify_ = true;
|
|
|
|
void SetUp() override
|
|
{
|
|
if constexpr(Rank == 4)
|
|
{
|
|
in_lengths_ = std::vector<std::vector<index_t>>{
|
|
{1, 2, 128, 1024}, {2, 4, 16, 8448}, {1, 128, 128, 64}};
|
|
}
|
|
}
|
|
|
|
void RunSingle(std::vector<index_t> in_length,
|
|
std::vector<index_t> reduce_dims,
|
|
AccDataType alpha,
|
|
AccDataType beta)
|
|
{
|
|
int init_method = 1; // integer value initialization
|
|
bool log = false;
|
|
std::vector<ck::index_t> strides; // intenionally empty, to get packed layout.
|
|
bool pass = false;
|
|
|
|
if constexpr(Rank == 3)
|
|
{
|
|
if(reduce_dims.size() == 1)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 1>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
else if(reduce_dims.size() == 2)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 2>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
else if(reduce_dims.size() == 3)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 3>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
}
|
|
else if constexpr(Rank == 4)
|
|
{
|
|
if(reduce_dims.size() == 1)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 1>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
else if(reduce_dims.size() == 2)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 2>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
else if(reduce_dims.size() == 3)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 3>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
else if(reduce_dims.size() == 4)
|
|
pass = ck::profiler::
|
|
profile_softmax_impl<InDataType, AccDataType, OutDataType, Rank, 4>(verify_,
|
|
init_method,
|
|
log,
|
|
bench_,
|
|
in_length,
|
|
strides,
|
|
reduce_dims,
|
|
alpha,
|
|
beta);
|
|
};
|
|
|
|
EXPECT_TRUE(pass);
|
|
}
|
|
|
|
void Run(std::vector<index_t> reduce_dims = {})
|
|
{
|
|
if(reduce_dims.empty())
|
|
{
|
|
reduce_dims.push_back(Rank - 1);
|
|
}
|
|
|
|
for(auto in_length : this->in_lengths_)
|
|
{
|
|
for(auto scale : this->scales_)
|
|
{
|
|
this->RunSingle(in_length, reduce_dims, scale[0], scale[1]);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
template <index_t Rank,
|
|
index_t NumReduceDim,
|
|
index_t BlockSize,
|
|
index_t MThreadClusterSize,
|
|
index_t KThreadClusterSize,
|
|
index_t MThreadSliceSize,
|
|
index_t KThreadSliceSize,
|
|
index_t InSrcVectorDim,
|
|
index_t InSrcVectorSize,
|
|
index_t OutDstVectorSize>
|
|
struct DeviceSoftmaxInstanceWrapper
|
|
{
|
|
using F16 = half_t;
|
|
using F32 = float;
|
|
using Pass = tensor_operation::element_wise::PassThrough;
|
|
|
|
using InDataType = F16;
|
|
using AccDataType = F32;
|
|
using OutDataType = F16;
|
|
using InElementOp = Pass;
|
|
using AccElementOp = Pass;
|
|
|
|
using DeviceSoftmaxInstance = tensor_operation::device::DeviceSoftmaxImpl<InDataType,
|
|
AccDataType,
|
|
OutDataType,
|
|
InElementOp,
|
|
AccElementOp,
|
|
Rank,
|
|
NumReduceDim,
|
|
BlockSize,
|
|
MThreadClusterSize,
|
|
KThreadClusterSize,
|
|
MThreadSliceSize,
|
|
KThreadSliceSize,
|
|
InSrcVectorDim,
|
|
InSrcVectorSize,
|
|
OutDstVectorSize>;
|
|
|
|
bool IsSupported(const std::vector<index_t> in_lengths,
|
|
const std::vector<index_t> in_strides,
|
|
const std::vector<index_t> reduce_dims) const
|
|
{
|
|
auto softmax = DeviceSoftmaxInstance{};
|
|
auto argument = softmax.MakeArgument(in_lengths,
|
|
in_strides,
|
|
reduce_dims,
|
|
1, // alpha
|
|
1, // beta
|
|
nullptr, // in_dev
|
|
nullptr, // in_out
|
|
Pass{}, // in elementwise op
|
|
Pass{}); // acc elementwise op
|
|
return softmax.IsSupportedArgument(argument);
|
|
}
|
|
};
|
|
|
|
} // namespace ck
|