reduction_operator_mapping.hpp Source File

reduction_operator_mapping.hpp Source File#

Composable Kernel: reduction_operator_mapping.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck/utility/reduction_operator.hpp"
#include "ck/utility/reduction_enums.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
// FIXME: can it be replaced with ck::Tuple?
#include <tuple>
 
namespace ck {
 
// The templated struct reduce_binary_operator maps the enum Ids of binary operators to their
// respective functor classes.
// The boolean member "indexable" are also provided in reduce_binary_operactor for
// easier checking by the upper-layer codes in the kernels.
 
template <ReduceTensorOp Op>
struct reduce_binary_operator;
 
template <>
struct reduce_binary_operator<ReduceTensorOp::ADD>
{
    using opType = reduce::Add;
 
    static constexpr bool indexable = false;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::MUL>
{
    using opType = reduce::Mul;
 
    static constexpr bool indexable = false;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::MIN>
{
    using opType = reduce::Min;
 
    static constexpr bool indexable = true;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::MAX>
{
    using opType = reduce::Max;
 
    static constexpr bool indexable = true;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::AMAX>
{
    using opType = reduce::AMax;
 
    static constexpr bool indexable = true;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::AVG>
{
    using opType = reduce::Add;
 
    static constexpr bool indexable = false;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::NORM1>
{
    using opType = reduce::Add;
 
    static constexpr bool indexable = false;
};
 
template <>
struct reduce_binary_operator<ReduceTensorOp::NORM2>
{
    using opType = reduce::Add;
 
    static constexpr bool indexable = false;
};
 
// The templated struct reduce_unary_operator maps the enum Ids of Reduce operators to two unary
// functor classes.
// The two unary functors are called before and afer the Reduction is executed respectively
template <ReduceTensorOp Op, bool IsFirstReduce, bool IsLastReduce>
struct reduce_unary_operator
{
    using InElementwiseOperation  = tensor_operation::element_wise::PassThrough;
    using AccElementwiseOperation = tensor_operation::element_wise::PassThrough;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
template <bool IsFirstReduce>
struct reduce_unary_operator<ReduceTensorOp::AVG, IsFirstReduce, true>
{
    using InElementwiseOperation  = tensor_operation::element_wise::PassThrough;
    using AccElementwiseOperation = tensor_operation::element_wise::UnaryDivide;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{reduceLength});
    };
};
 
template <bool IsLastReduce>
struct reduce_unary_operator<ReduceTensorOp::NORM1, true, IsLastReduce>
{
    using InElementwiseOperation  = tensor_operation::element_wise::UnaryAbs;
    using AccElementwiseOperation = tensor_operation::element_wise::PassThrough;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
template <bool IsLastReduce>
struct reduce_unary_operator<ReduceTensorOp::AMAX, true, IsLastReduce>
{
    using InElementwiseOperation  = tensor_operation::element_wise::UnaryAbs;
    using AccElementwiseOperation = tensor_operation::element_wise::PassThrough;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
template <>
struct reduce_unary_operator<ReduceTensorOp::NORM2, true, false>
{
    using InElementwiseOperation  = tensor_operation::element_wise::UnarySquare;
    using AccElementwiseOperation = tensor_operation::element_wise::PassThrough;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
template <>
struct reduce_unary_operator<ReduceTensorOp::NORM2, true, true>
{
    using InElementwiseOperation  = tensor_operation::element_wise::UnarySquare;
    using AccElementwiseOperation = tensor_operation::element_wise::UnarySqrt;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
template <>
struct reduce_unary_operator<ReduceTensorOp::NORM2, false, true>
{
    using InElementwiseOperation  = tensor_operation::element_wise::PassThrough;
    using AccElementwiseOperation = tensor_operation::element_wise::UnarySqrt;
 
    static std::tuple<InElementwiseOperation, AccElementwiseOperation>
    GetElementwiseOperator(int32_t reduceLength)
    {
        (void)reduceLength;
        return std::make_tuple(InElementwiseOperation{}, AccElementwiseOperation{});
    };
};
 
} // namespace ck