mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-17 11:30:02 +00:00
Update to the Reduction API and instances (#476)
* Simplify the macros for declaring and defining the add_device_reduce_instance_xxxx() instances
* Change the types of lengths and strides from std::vector to std::array for the reduction device interfaces
* Remove DeviceSoftmaxImpl's depending on DeviceReduceMultiblock
* Split the cpp and hpp files for reduction instances to enable more parallel compiling
* Remove the using of macros for declaring reduction instances and instance references
* Update to add_device_reduce_instance_xxxx templated functions
* Use ReduceOperation+InElementwiseOp+AccElementwiseOp to repace the ReduceOpId in defining add_reduce_instance_xxxx() templates
* Change return format
[ROCm/composable_kernel commit: dda3a0a10b]
This commit is contained in:
@@ -140,6 +140,10 @@ bool reduce_blockwise_test(bool do_verification,
|
||||
if(ShapeType::Rank_ != inLengths.size() || ShapeType::NumReduceDim_ != reduceDims.size())
|
||||
return;
|
||||
|
||||
std::array<int, ShapeType::NumReduceDim_> arrReduceDims;
|
||||
|
||||
std::copy(reduceDims.begin(), reduceDims.end(), arrReduceDims.begin());
|
||||
|
||||
result = reduce_blockwise_impl<InOutDataType,
|
||||
AccDataType,
|
||||
ReduceOpId,
|
||||
@@ -147,7 +151,7 @@ bool reduce_blockwise_test(bool do_verification,
|
||||
ShapeType::NumReduceDim_,
|
||||
PropagateNan,
|
||||
OutputIndex>(
|
||||
do_verification, init_method, time_kernel, inLengths, reduceDims, alpha, beta);
|
||||
do_verification, init_method, time_kernel, inLengths, arrReduceDims, alpha, beta);
|
||||
|
||||
matched = true;
|
||||
});
|
||||
|
||||
@@ -30,7 +30,7 @@ int reduce_blockwise_impl(bool do_verification,
|
||||
int init_method,
|
||||
bool time_kernel,
|
||||
const std::vector<size_t>& inLengths,
|
||||
const std::vector<int>& reduceDims,
|
||||
const std::array<int, NumReduceDim>& reduceDims,
|
||||
float alpha,
|
||||
float beta)
|
||||
|
||||
@@ -38,6 +38,8 @@ int reduce_blockwise_impl(bool do_verification,
|
||||
using namespace ck;
|
||||
using namespace ck::tensor_operation::device;
|
||||
|
||||
constexpr index_t NumOutDim = (Rank - NumReduceDim == 0) ? 1 : Rank - NumReduceDim;
|
||||
|
||||
constexpr bool op_support_indices =
|
||||
(ReduceOpId == ReduceTensorOp::MIN || ReduceOpId == ReduceTensorOp::MAX ||
|
||||
ReduceOpId == ReduceTensorOp::AMAX);
|
||||
@@ -143,7 +145,7 @@ int reduce_blockwise_impl(bool do_verification,
|
||||
|
||||
std::vector<size_t> outLengths;
|
||||
|
||||
std::vector<int> invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
|
||||
auto invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
|
||||
|
||||
if(invariantDims.empty())
|
||||
outLengths.push_back(1);
|
||||
@@ -256,22 +258,22 @@ int reduce_blockwise_impl(bool do_verification,
|
||||
acc_elementwise_op);
|
||||
};
|
||||
|
||||
std::vector<ck::index_t> i_inLengths;
|
||||
std::vector<ck::index_t> i_inStrides;
|
||||
std::vector<ck::index_t> i_outLengths;
|
||||
std::vector<ck::index_t> i_outStrides;
|
||||
std::array<index_t, Rank> arrInLengths;
|
||||
std::array<index_t, Rank> arrInStrides;
|
||||
std::array<index_t, NumOutDim> arrOutLengths;
|
||||
std::array<index_t, NumOutDim> arrOutStrides;
|
||||
|
||||
i_inLengths.assign(inLengths.begin(), inLengths.end());
|
||||
i_inStrides.assign(inStrides.begin(), inStrides.end());
|
||||
i_outLengths.assign(outLengths.begin(), outLengths.end());
|
||||
i_outStrides.assign(outStrides.begin(), outStrides.end());
|
||||
std::copy(inLengths.begin(), inLengths.end(), arrInLengths.begin());
|
||||
std::copy(inStrides.begin(), inStrides.end(), arrInStrides.begin());
|
||||
std::copy(outLengths.begin(), outLengths.end(), arrOutLengths.begin());
|
||||
std::copy(outStrides.begin(), outStrides.end(), arrOutStrides.begin());
|
||||
|
||||
auto reduce = DeviceReduceInstance{};
|
||||
|
||||
auto argument_ptr = reduce.MakeArgumentPointer(i_inLengths,
|
||||
i_inStrides,
|
||||
i_outLengths,
|
||||
i_outStrides,
|
||||
auto argument_ptr = reduce.MakeArgumentPointer(arrInLengths,
|
||||
arrInStrides,
|
||||
arrOutLengths,
|
||||
arrOutStrides,
|
||||
reduceDims,
|
||||
alpha,
|
||||
beta,
|
||||
|
||||
@@ -90,15 +90,15 @@ static bool time_kernel;
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
// used by the device reduction
|
||||
const std::vector<int> reduceDims_1 = {4};
|
||||
const std::vector<int> invariantDims_1 = {0, 1, 2, 3};
|
||||
const std::array<int, 1> reduceDims_1 = {4};
|
||||
// const std::array<int, 4> invariantDims_1 = {0, 1, 2, 3};
|
||||
|
||||
const std::vector<int> reduceDims_2 = {3};
|
||||
const std::vector<int> invariantDims_2 = {0, 1, 2};
|
||||
const std::array<int, 1> reduceDims_2 = {3};
|
||||
// const std::array<int, 3> invariantDims_2 = {0, 1, 2};
|
||||
|
||||
// used by the host reduction
|
||||
const std::vector<int> reduceDims = {3, 4};
|
||||
const std::vector<int> invariantDims = {0, 1, 2};
|
||||
const std::array<int, 2> reduceDims = {3, 4};
|
||||
const std::array<int, 3> invariantDims = {0, 1, 2};
|
||||
|
||||
const std::vector<size_t> inLengths_1 = {64, 320, 80, 4, 128};
|
||||
|
||||
@@ -214,26 +214,26 @@ int main(int argc, char* argv[])
|
||||
acc_elementwise_op);
|
||||
};
|
||||
|
||||
std::vector<ck::index_t> i_inLengths_1;
|
||||
std::vector<ck::index_t> i_inStrides_1;
|
||||
std::vector<ck::index_t> i_inLengths_2;
|
||||
std::vector<ck::index_t> i_inStrides_2;
|
||||
std::vector<ck::index_t> i_outLengths;
|
||||
std::vector<ck::index_t> i_outStrides;
|
||||
std::array<index_t, 5> arrInLengths_1;
|
||||
std::array<index_t, 5> arrInStrides_1;
|
||||
std::array<index_t, 4> arrInLengths_2;
|
||||
std::array<index_t, 4> arrInStrides_2;
|
||||
std::array<index_t, 3> arrOutLengths;
|
||||
std::array<index_t, 3> arrOutStrides;
|
||||
|
||||
i_inLengths_1.assign(inLengths_1.begin(), inLengths_1.end());
|
||||
i_inStrides_1.assign(inStrides_1.begin(), inStrides_1.end());
|
||||
i_inLengths_2.assign(inLengths_2.begin(), inLengths_2.end());
|
||||
i_inStrides_2.assign(inStrides_2.begin(), inStrides_2.end());
|
||||
i_outLengths.assign(outLengths.begin(), outLengths.end());
|
||||
i_outStrides.assign(outStrides.begin(), outStrides.end());
|
||||
std::copy(inLengths_1.begin(), inLengths_1.end(), arrInLengths_1.begin());
|
||||
std::copy(inStrides_1.begin(), inStrides_1.end(), arrInStrides_1.begin());
|
||||
std::copy(inLengths_2.begin(), inLengths_2.end(), arrInLengths_2.begin());
|
||||
std::copy(inStrides_2.begin(), inStrides_2.end(), arrInStrides_2.begin());
|
||||
std::copy(outLengths.begin(), outLengths.end(), arrOutLengths.begin());
|
||||
std::copy(outStrides.begin(), outStrides.end(), arrOutStrides.begin());
|
||||
|
||||
auto reduce_1 = DeviceReduceInstance_1{};
|
||||
|
||||
auto argument_ptr_1 = reduce_1.MakeArgumentPointer(i_inLengths_1,
|
||||
i_inStrides_1,
|
||||
i_inLengths_2,
|
||||
i_inStrides_2,
|
||||
auto argument_ptr_1 = reduce_1.MakeArgumentPointer(arrInLengths_1,
|
||||
arrInStrides_1,
|
||||
arrInLengths_2,
|
||||
arrInStrides_2,
|
||||
reduceDims_1,
|
||||
1.0f,
|
||||
0.0f,
|
||||
@@ -255,10 +255,10 @@ int main(int argc, char* argv[])
|
||||
|
||||
auto reduce_2 = DeviceReduceInstance_2{};
|
||||
|
||||
auto argument_ptr_2 = reduce_2.MakeArgumentPointer(i_inLengths_2,
|
||||
i_inStrides_2,
|
||||
i_outLengths,
|
||||
i_outStrides,
|
||||
auto argument_ptr_2 = reduce_2.MakeArgumentPointer(arrInLengths_2,
|
||||
arrInStrides_2,
|
||||
arrOutLengths,
|
||||
arrOutStrides,
|
||||
reduceDims_2,
|
||||
alpha,
|
||||
beta,
|
||||
|
||||
@@ -5,11 +5,10 @@
|
||||
|
||||
#include "ck/ck.hpp"
|
||||
|
||||
template <ck::index_t Rank, ck::index_t NumReduceDim>
|
||||
std::vector<int> get_invariant_dims(const std::vector<int>& reduceDims)
|
||||
template <int Rank, int NumReduceDim>
|
||||
static inline std::array<int, Rank - NumReduceDim>
|
||||
get_invariant_dims(const std::array<int, NumReduceDim>& reduceDims)
|
||||
{
|
||||
assert(NumReduceDim == reduceDims.size());
|
||||
|
||||
int reduceFlag = 0;
|
||||
|
||||
// flag the bits for the reduceDims
|
||||
@@ -18,13 +17,15 @@ std::vector<int> get_invariant_dims(const std::vector<int>& reduceDims)
|
||||
reduceFlag |= 1 << reduceDims[i];
|
||||
};
|
||||
|
||||
std::vector<int> invariantDims;
|
||||
std::array<int, Rank - NumReduceDim> invariantDims;
|
||||
|
||||
// collect invariant dimensions
|
||||
int dim = 0;
|
||||
for(int i = 0; i < Rank; i++)
|
||||
if((reduceFlag & (1 << i)) == 0)
|
||||
{
|
||||
invariantDims.push_back(i);
|
||||
invariantDims[dim] = i;
|
||||
dim++;
|
||||
};
|
||||
|
||||
return invariantDims;
|
||||
|
||||
@@ -138,13 +138,17 @@ bool reduce_multiblock_atomic_add_test(bool do_verification,
|
||||
if(ShapeType::Rank_ != inLengths.size() || ShapeType::NumReduceDim_ != reduceDims.size())
|
||||
return;
|
||||
|
||||
std::array<int, ShapeType::NumReduceDim_> a_reduceDims;
|
||||
|
||||
std::copy(reduceDims.begin(), reduceDims.end(), a_reduceDims.begin());
|
||||
|
||||
result = reduce_multiblock_atomic_add_impl<InOutDataType,
|
||||
AccDataType,
|
||||
ReduceOpId,
|
||||
ShapeType::Rank_,
|
||||
ShapeType::NumReduceDim_,
|
||||
PropagateNan>(
|
||||
do_verification, init_method, time_kernel, inLengths, reduceDims, alpha, beta);
|
||||
do_verification, init_method, time_kernel, inLengths, a_reduceDims, alpha, beta);
|
||||
|
||||
matched = true;
|
||||
});
|
||||
|
||||
@@ -29,7 +29,7 @@ int reduce_multiblock_atomic_add_impl(bool do_verification,
|
||||
int init_method,
|
||||
bool time_kernel,
|
||||
const std::vector<size_t>& inLengths,
|
||||
const std::vector<int>& reduceDims,
|
||||
const std::array<int, NumReduceDim>& reduceDims,
|
||||
float alpha,
|
||||
float beta)
|
||||
|
||||
@@ -37,6 +37,8 @@ int reduce_multiblock_atomic_add_impl(bool do_verification,
|
||||
using namespace ck;
|
||||
using namespace ck::tensor_operation::device;
|
||||
|
||||
constexpr index_t NumOutDim = (Rank - NumReduceDim == 0) ? 1 : Rank - NumReduceDim;
|
||||
|
||||
constexpr bool op_support_atomic_add =
|
||||
(ReduceOpId == ReduceTensorOp::ADD || ReduceOpId == ReduceTensorOp::AVG);
|
||||
|
||||
@@ -84,7 +86,7 @@ int reduce_multiblock_atomic_add_impl(bool do_verification,
|
||||
|
||||
std::vector<size_t> outLengths;
|
||||
|
||||
std::vector<int> invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
|
||||
auto invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
|
||||
|
||||
if(invariantDims.empty())
|
||||
outLengths.push_back(1);
|
||||
@@ -169,22 +171,22 @@ int reduce_multiblock_atomic_add_impl(bool do_verification,
|
||||
acc_elementwise_op);
|
||||
};
|
||||
|
||||
std::vector<ck::index_t> i_inLengths;
|
||||
std::vector<ck::index_t> i_inStrides;
|
||||
std::vector<ck::index_t> i_outLengths;
|
||||
std::vector<ck::index_t> i_outStrides;
|
||||
std::array<index_t, Rank> arrInLengths;
|
||||
std::array<index_t, Rank> arrInStrides;
|
||||
std::array<index_t, NumOutDim> arrOutLengths;
|
||||
std::array<index_t, NumOutDim> arrOutStrides;
|
||||
|
||||
i_inLengths.assign(inLengths.begin(), inLengths.end());
|
||||
i_inStrides.assign(inStrides.begin(), inStrides.end());
|
||||
i_outLengths.assign(outLengths.begin(), outLengths.end());
|
||||
i_outStrides.assign(outStrides.begin(), outStrides.end());
|
||||
std::copy(inLengths.begin(), inLengths.end(), arrInLengths.begin());
|
||||
std::copy(inStrides.begin(), inStrides.end(), arrInStrides.begin());
|
||||
std::copy(outLengths.begin(), outLengths.end(), arrOutLengths.begin());
|
||||
std::copy(outStrides.begin(), outStrides.end(), arrOutStrides.begin());
|
||||
|
||||
auto reduce = DeviceReduceInstance{};
|
||||
|
||||
auto argument_ptr = reduce.MakeArgumentPointer(i_inLengths,
|
||||
i_inStrides,
|
||||
i_outLengths,
|
||||
i_outStrides,
|
||||
auto argument_ptr = reduce.MakeArgumentPointer(arrInLengths,
|
||||
arrInStrides,
|
||||
arrOutLengths,
|
||||
arrOutStrides,
|
||||
reduceDims,
|
||||
alpha,
|
||||
beta,
|
||||
|
||||
Reference in New Issue
Block a user