/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp Source File

Go to the documentation of this 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