/* ****************************************************************************** * * * This program and the accompanying materials are made available under the * terms of the Apache License, Version 2.0 which is available at * https://www.apache.org/licenses/LICENSE-2.0. * * See the NOTICE file distributed with this work for additional * information regarding copyright ownership. * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the * License for the specific language governing permissions and limitations * under the License. * * SPDX-License-Identifier: Apache-2.0 ******************************************************************************/ #pragma once #ifndef OP_MACROS_REDUCE_H_ #define OP_MACROS_REDUCE_H_ #include "op_types.h" namespace simdOps { // ============================================================================= // BASIC REDUCE OPERATION MACROS // ============================================================================= /** * @brief Declares a basic reduce operation with proper SIMD handling */ #define DECLARE_REDUCE_OP(OP_NAME, REDUCE_TYPE_VAL, STARTING_VAL, MERGE_OP, UPDATE_OP, POST_PROCESS) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static X op_logic(X d1, X* extraParams) { return d1; } \ SD_HOST_DEVICE SD_INLINE static X merge_logic(X old, X opOutput, X* extraParams) { return MERGE_OP; } \ SD_HOST_DEVICE SD_INLINE static X update_logic(X old, X opOutput, X* extraParams) { return UPDATE_OP; } \ SD_HOST_DEVICE SD_INLINE static X postProcess_logic(X reduction, sd::LongType n, X* extraParams) { return POST_PROCESS; } \ static SD_HOST_DEVICE X op_simd(X d1, X* extraParams) { return op_logic(d1, extraParams); } \ static SD_HOST_DEVICE X merge_simd(X old, X opOutput, X* extraParams) { return merge_logic(old, opOutput, extraParams); } \ static SD_HOST_DEVICE X update_simd(X old, X opOutput, X* extraParams) { return update_logic(old, opOutput, extraParams); } \ static SD_HOST_DEVICE X postProcess_simd(X reduction, sd::LongType n, X* extraParams) { \ return postProcess_logic(reduction, n, extraParams); \ } \ \ public: \ no_op_exec_special_accumulation_same no_op_exec_special_accumulation_same_cuda \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ using InterType = typename AggregateType::type; \ \ static SD_HOST_DEVICE X startingValue(const X* input) { return STARTING_VAL; } \ static SD_HOST_DEVICE X merge(X old, X opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParams); \ else \ return merge_simd(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE X update(X old, X opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParams); \ else \ return update_simd(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE X op(X d1, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) return op_logic(d1, extraParams); \ else return op_simd(d1, extraParams); \ } \ static SD_HOST_DEVICE X postProcess(X reduction, sd::LongType n, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return postProcess_logic(reduction, n, extraParams); \ else \ return postProcess_simd(reduction, n, extraParams); \ } \ }; // ============================================================================= // FLOAT REDUCE OPERATION MACROS // ============================================================================= #define DECLARE_REDUCE_FLOAT_OP(OP_NAME, REDUCE_TYPE_VAL, STARTING_VAL, MERGE_OP, UPDATE_OP, POST_PROCESS) \ template \ class OP_NAME { \ public: \ using InterType = typename AggregateType::type; \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ private: \ SD_HOST_DEVICE SD_INLINE static Z postProcess_logic(InterType reduction, sd::LongType n, Z* extraParams) { \ return POST_PROCESS; \ } \ SD_HOST_DEVICE SD_INLINE static InterType merge_logic(InterType old, InterType opOutput, Z* extraParams) { \ return MERGE_OP; \ } \ SD_HOST_DEVICE SD_INLINE static InterType update_logic(InterType old, InterType opOutput, Z* extraParams) { \ return UPDATE_OP; \ } \ SD_HOST_DEVICE SD_INLINE static InterType op_logic(X d1, Z* extraParams) { \ return static_cast(d1); \ } \ static SD_HOST_DEVICE Z postProcess_simd(InterType reduction, sd::LongType n, Z* extraParams) { \ return postProcess_logic(reduction, n, extraParams); \ } \ static SD_HOST_DEVICE InterType merge_simd(InterType old, InterType opOutput, Z* extraParams) { \ return merge_logic(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE InterType update_simd(InterType old, InterType opOutput, Z* extraParams) { \ return update_logic(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE InterType op_simd(X d1, Z* extraParams) { \ return op_logic(d1, extraParams); \ } \ \ public: \ no_op_exec_special_accumulation no_op_exec_special_accumulation_cuda; \ static SD_HOST_DEVICE X startingValue(const X* input) { return STARTING_VAL; } \ static SD_HOST_DEVICE InterType op(X d1, Z* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic(d1, extraParams); \ else return op_simd(d1, extraParams); \ } \ static SD_HOST_DEVICE InterType merge(InterType old, InterType opOutput, Z* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParams); \ else return merge_simd(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE InterType update(InterType old, InterType opOutput, Z* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParams); \ else return update_simd(old, opOutput, extraParams); \ } \ static SD_HOST_DEVICE Z postProcess(InterType reduction, sd::LongType n, Z* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return postProcess_logic(reduction, n, extraParams); \ else return postProcess_simd(reduction, n, extraParams); \ } \ }; /** * @brief Fixed DECLARE_REDUCE3_OP_WITH_BOOL_SUPPORT macro * Uses conditional compilation to include opAtomic only when atomics are available */ /** * @brief DECLARE_REDUCE3_OP_WITH_BOOL_SUPPORT macro with conditional opAtomic * Two versions: one with atomics for CUDA, one without for host */ #ifdef __CUDACC__ #define DECLARE_REDUCE3_OP_WITH_BOOL_SUPPORT(OP_NAME, BOOL_LOGIC, NORMAL_LOGIC, EXTRA_PARAMS_LEN, STARTING_VAL, POST_PROCESS) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static Y op_logic(X d1, X d2, Y* extraParamsRef) { \ Y* extraParams = extraParamsRef; \ if constexpr (std::is_same_v) { \ BOOL_LOGIC \ } else { \ NORMAL_LOGIC \ } \ } \ \ SD_HOST_DEVICE SD_INLINE static Y startingValue_logic(const X* input) { \ return static_cast(STARTING_VAL); \ } \ \ SD_HOST_DEVICE SD_INLINE static Y postProcess_logic(Y reduction, sd::LongType n, Y* extraParamsRef) { \ Y* extraParams = extraParamsRef; \ return POST_PROCESS; \ } \ \ SD_HOST_DEVICE SD_INLINE static Y update_logic(Y old, Y opOutput, Y* extraParamsRef) { \ return old + opOutput; \ } \ \ SD_HOST_DEVICE SD_INLINE static Y merge_logic(Y old, Y opOutput, Y* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd(TX d1, TX d2, TY* extraParamsRef) { \ return op_logic(d1, d2, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd(TX d1, TX d2, TY* extraParamsRef) { \ return op_logic(d1, d2, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe startingValue_simd(const TX* input) { \ return startingValue_logic(input); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe startingValue_simd(const TX* input) { \ return startingValue_logic(input); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe postProcess_simd(TY reduction, sd::LongType n, TY* extraParamsRef) { \ return postProcess_logic(reduction, n, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe postProcess_simd(TY reduction, sd::LongType n, TY* extraParamsRef) { \ return postProcess_logic(reduction, n, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe update_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe update_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe merge_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return merge_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe merge_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return merge_logic(old, opOutput, extraParamsRef); \ } \ \ public: \ static const int extraParamsLen = EXTRA_PARAMS_LEN; \ static SD_HOST_DEVICE X *generateExtraParams() { return nullptr; } \ static SD_HOST_DEVICE void finalizeExtraParams(X *extraParamsRef) {} \ static SD_HOST_DEVICE void aggregateExtraParams(Y *extraParamsTotal, Y *extraParamsLocal) { \ for (int i = 0; i < EXTRA_PARAMS_LEN; i++) { \ extraParamsTotal[i] += extraParamsLocal[i]; \ } \ } \ \ static SD_HOST_DEVICE Y startingValue(const X* input) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return startingValue_logic(input); \ else \ return startingValue_simd(input); \ } \ \ static SD_HOST_DEVICE Y postProcess(Y reduction, sd::LongType n, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return postProcess_logic(reduction, n, extraParamsRef); \ else \ return postProcess_simd(reduction, n, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y op(X d1, X d2, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic(d1, d2, extraParamsRef); \ else \ return op_simd(d1, d2, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y update(Y old, Y opOutput, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParamsRef); \ else \ return update_simd(old, opOutput, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y merge(Y old, Y opOutput, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParamsRef); \ else \ return merge_simd(old, opOutput, extraParamsRef); \ } \ \ static SD_DEVICE Y opAtomic(X d1, X d2, Y* extraParamsRef) { \ if constexpr (std::is_same_v) { \ Y val1 = static_cast(static_cast(d1) * static_cast(d1)); \ Y val2 = static_cast(static_cast(d2) * static_cast(d2)); \ sd::math::atomics::sd_atomicAdd(&extraParamsRef[0], val1); \ sd::math::atomics::sd_atomicAdd(&extraParamsRef[1], val2); \ return static_cast(static_cast(d1) * static_cast(d2)); \ } else { \ Y val1 = static_cast(sd::math::sd_abs(d1) * sd::math::sd_abs(d1)); \ Y val2 = static_cast(sd::math::sd_abs(d2) * sd::math::sd_abs(d2)); \ sd::math::atomics::sd_atomicAdd(&extraParamsRef[0], val1); \ sd::math::atomics::sd_atomicAdd(&extraParamsRef[1], val2); \ return static_cast(d1 * d2); \ } \ } \ }; #else #define DECLARE_REDUCE3_OP_WITH_BOOL_SUPPORT(OP_NAME, BOOL_LOGIC, NORMAL_LOGIC, EXTRA_PARAMS_LEN, STARTING_VAL, POST_PROCESS) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static Y op_logic(X d1, X d2, Y* extraParamsRef) { \ Y* extraParams = extraParamsRef; \ if constexpr (std::is_same_v) { \ BOOL_LOGIC \ } else { \ NORMAL_LOGIC \ } \ } \ \ SD_HOST_DEVICE SD_INLINE static Y startingValue_logic(const X* input) { \ return static_cast(STARTING_VAL); \ } \ \ SD_HOST_DEVICE SD_INLINE static Y postProcess_logic(Y reduction, sd::LongType n, Y* extraParamsRef) { \ Y* extraParams = extraParamsRef; \ return POST_PROCESS; \ } \ \ SD_HOST_DEVICE SD_INLINE static Y update_logic(Y old, Y opOutput, Y* extraParamsRef) { \ return old + opOutput; \ } \ \ SD_HOST_DEVICE SD_INLINE static Y merge_logic(Y old, Y opOutput, Y* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd(TX d1, TX d2, TY* extraParamsRef) { \ return op_logic(d1, d2, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd(TX d1, TX d2, TY* extraParamsRef) { \ return op_logic(d1, d2, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe startingValue_simd(const TX* input) { \ return startingValue_logic(input); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe startingValue_simd(const TX* input) { \ return startingValue_logic(input); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe postProcess_simd(TY reduction, sd::LongType n, TY* extraParamsRef) { \ return postProcess_logic(reduction, n, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe postProcess_simd(TY reduction, sd::LongType n, TY* extraParamsRef) { \ return postProcess_logic(reduction, n, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe update_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe update_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return update_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe merge_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return merge_logic(old, opOutput, extraParamsRef); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe merge_simd(TY old, TY opOutput, TY* extraParamsRef) { \ return merge_logic(old, opOutput, extraParamsRef); \ } \ \ public: \ static const int extraParamsLen = EXTRA_PARAMS_LEN; \ static SD_HOST_DEVICE X *generateExtraParams() { return nullptr; } \ static SD_HOST_DEVICE void finalizeExtraParams(X *extraParamsRef) {} \ static SD_HOST_DEVICE void aggregateExtraParams(Y *extraParamsTotal, Y *extraParamsLocal) { \ for (int i = 0; i < EXTRA_PARAMS_LEN; i++) { \ extraParamsTotal[i] += extraParamsLocal[i]; \ } \ } \ \ static SD_HOST_DEVICE Y startingValue(const X* input) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return startingValue_logic(input); \ else \ return startingValue_simd(input); \ } \ \ static SD_HOST_DEVICE Y postProcess(Y reduction, sd::LongType n, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return postProcess_logic(reduction, n, extraParamsRef); \ else \ return postProcess_simd(reduction, n, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y op(X d1, X d2, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic(d1, d2, extraParamsRef); \ else \ return op_simd(d1, d2, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y update(Y old, Y opOutput, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParamsRef); \ else \ return update_simd(old, opOutput, extraParamsRef); \ } \ \ static SD_HOST_DEVICE Y merge(Y old, Y opOutput, Y* extraParamsRef) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParamsRef); \ else \ return merge_simd(old, opOutput, extraParamsRef); \ } \ }; #endif #define DECLARE_REDUCE_SIMD_SAFE_OP(OP_NAME, OPERATION) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static Z op_logic(X d1, Z* params) { \ OPERATION \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd(TX d1, TZ* params) { \ return op_logic(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd(TX d1, TZ* params) { \ return op_logic(d1, params); \ } \ \ public: \ no_op_exec_special_bool no_op_exec_special_bool_cuda \ no_op_exec_special_accumulation no_op_exec_special_accumulation_cuda \ \ using InterType = typename AggregateType::type; \ \ static SD_HOST_DEVICE Z op(X d1, Z* params) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic(d1, params); \ else \ return op_simd(d1, params); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if, Z>::type op(X d1, X* extraParams) { \ return op(d1, reinterpret_cast(extraParams)); \ } \ \ static SD_HOST_DEVICE X startingValue(const X* input) { return static_cast(0); } \ \ static SD_HOST_DEVICE InterType merge(InterType old, InterType opOutput, Z* extraParams) { return opOutput + old; } \ \ static SD_HOST_DEVICE InterType update(InterType old, InterType opOutput, Z* extraParams) { return opOutput + old; } \ \ template \ static SD_HOST_DEVICE typename std::enable_if, InterType>::type update(InterType old, InterType opOutput, X* extraParams) { \ return update(old, opOutput, reinterpret_cast(extraParams)); \ } \ \ static SD_HOST_DEVICE Z postProcess(InterType reduction, sd::LongType n, Z* extraParams) { return static_cast(reduction); } \ \ template \ static SD_HOST_DEVICE typename std::enable_if, Z>::type postProcess(Z reduction, sd::LongType n, X* extraParams) { \ return postProcess(reduction, n, reinterpret_cast(extraParams)); \ } \ }; #define DECLARE_COMPLEX_ACCUMULATION_SIMD_SAFE_OP(OP_NAME, OP_LOGIC, OP_BINARY_LOGIC, OP_BINARY_NO_PARAMS_LOGIC, REDUCE_TYPE_VAL, STARTING_VAL, MERGE_OP, UPDATE_OP, POST_PROCESS) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static X op_logic(X d1, X* params) { \ OP_LOGIC \ } \ \ SD_HOST_DEVICE SD_INLINE static X op_binary_logic(X d1, X d2, X* params) { \ OP_BINARY_LOGIC \ } \ \ SD_HOST_DEVICE SD_INLINE static X op_binary_no_params_logic(X d1, X d2) { \ OP_BINARY_NO_PARAMS_LOGIC \ } \ \ SD_HOST_DEVICE SD_INLINE static X merge_logic(X old, X opOutput, X* extraParams) { \ return MERGE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static X update_logic(X old, X opOutput, X* extraParams) { \ return UPDATE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static X postProcess_logic(X reduction, sd::LongType n, X* extraParams) { \ return POST_PROCESS; \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd(TX d1, TX* params) { \ return op_logic(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd(TX d1, TX* params) { \ return op_logic(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_binary_simd(TX d1, TX d2, TX* params) { \ return op_binary_logic(d1, d2, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_binary_simd(TX d1, TX d2, TX* params) { \ return op_binary_logic(d1, d2, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_binary_no_params_simd(TX d1, TX d2) { \ return op_binary_no_params_logic(d1, d2); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_binary_no_params_simd(TX d1, TX d2) { \ return op_binary_no_params_logic(d1, d2); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe merge_simd(TX old, TX opOutput, TX* extraParams) { \ return merge_logic(old, opOutput, extraParams); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe merge_simd(TX old, TX opOutput, TX* extraParams) { \ return merge_logic(old, opOutput, extraParams); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe update_simd(TX old, TX opOutput, TX* extraParams) { \ return update_logic(old, opOutput, extraParams); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe update_simd(TX old, TX opOutput, TX* extraParams) { \ return update_logic(old, opOutput, extraParams); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe postProcess_simd(TX reduction, sd::LongType n, TX* extraParams) { \ return postProcess_logic(reduction, n, extraParams); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe postProcess_simd(TX reduction, sd::LongType n, TX* extraParams) { \ return postProcess_logic(reduction, n, extraParams); \ } \ \ public: \ no_op_exec_special_accumulation_same no_op_exec_special_accumulation_same_cuda \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ using InterType = typename AggregateType::type; \ \ static SD_HOST_DEVICE X startingValue(const X* input) { return STARTING_VAL; } \ \ static SD_HOST_DEVICE X op(X d1, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return op_logic(d1, extraParams); \ else \ return op_simd(d1, extraParams); \ } \ \ static SD_HOST_DEVICE X op(X d1, X d2, X* params) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return op_binary_logic(d1, d2, params); \ else \ return op_binary_simd(d1, d2, params); \ } \ \ static SD_HOST_DEVICE X op(X d1, X d2) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return op_binary_no_params_logic(d1, d2); \ else \ return op_binary_no_params_simd(d1, d2); \ } \ \ static SD_HOST_DEVICE X merge(X old, X opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParams); \ else \ return merge_simd(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE X update(X old, X opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParams); \ else \ return update_simd(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE X postProcess(X reduction, sd::LongType n, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return postProcess_logic(reduction, n, extraParams); \ else \ return postProcess_simd(reduction, n, extraParams); \ } \ }; #define DECLARE_SIMPLE_REDUCTION_OP(OP_NAME, REDUCE_TYPE_VAL, STARTING_VAL, OP_LOGIC, MERGE_LOGIC, UPDATE_LOGIC, POST_PROCESS_LOGIC) \ template \ class OP_NAME { \ public: \ static const bool requiresSpecialAccumulation = false; \ \ static void execSpecial(const X *x, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfoBuffer, sd::LongType *dimension, sd::LongType dimensionLength, \ const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {} \ \ static void execSpecialCuda(const X *dx, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, \ Z *reductionBuffer, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets) {} \ \ using InterType = typename AggregateType::type; \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ SD_OP_DEF static Z startingValue(const X *input) { \ return STARTING_VAL; \ } \ \ /* CORE FUNCTIONS - Now using Z* parameters for consistency */ \ SD_OP_DEF static InterType merge(InterType old, InterType opOutput, Z *extraParams) { \ return MERGE_LOGIC; \ } \ \ SD_OP_DEF static InterType update(InterType old, InterType opOutput, Z *extraParams) { \ return UPDATE_LOGIC; \ } \ \ SD_OP_DEF static InterType op(X d1, Z *extraParams) { \ return OP_LOGIC; \ } \ \ SD_OP_DEF static Z postProcess(InterType reduction, sd::LongType n, Z *extraParams) { \ return POST_PROCESS_LOGIC; \ } \ \ /* TEMPLATE OVERLOADS FOR DIFFERENT PARAMETER TYPES */ \ \ /* Overloads for X* parameters */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, InterType> \ merge(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return merge(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return merge(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, InterType> \ update(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return update(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return update(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, InterType> \ op(X d1, ParamType *extraParams) { \ if (extraParams == nullptr) return op(d1, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return op(d1, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, Z> \ postProcess(InterType reduction, sd::LongType n, ParamType *extraParams) { \ if (extraParams == nullptr) return postProcess(reduction, n, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return postProcess(reduction, n, convertedParams); \ } \ \ /* Overloads for sd::LongType* parameters - only when different from X */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ merge(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return merge(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return merge(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ update(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return update(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return update(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ op(X d1, ParamType *extraParams) { \ if (extraParams == nullptr) return op(d1, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return op(d1, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, Z> \ postProcess(InterType reduction, sd::LongType n, ParamType *extraParams) { \ if (extraParams == nullptr) return postProcess(reduction, n, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return postProcess(reduction, n, convertedParams); \ } \ \ /* Overloads for floating point types (like float* when Z=bfloat16) */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ merge(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return merge(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return merge(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ update(InterType old, InterType opOutput, ParamType *extraParams) { \ if (extraParams == nullptr) return update(old, opOutput, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return update(old, opOutput, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, InterType> \ op(X d1, ParamType *extraParams) { \ if (extraParams == nullptr) return op(d1, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return op(d1, convertedParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v && !std::is_same_v, Z> \ postProcess(InterType reduction, sd::LongType n, ParamType *extraParams) { \ if (extraParams == nullptr) return postProcess(reduction, n, static_cast(nullptr)); \ Z convertedParams[8]; \ for (int i = 0; i < 8; ++i) convertedParams[i] = static_cast(extraParams[i]); \ return postProcess(reduction, n, convertedParams); \ } \ }; #define DECLARE_ACCUMULATION_SIMD_SAFE_OP(OP_NAME, OPERATION, REDUCE_TYPE_VAL, STARTING_VAL, MERGE_OP, UPDATE_OP, POST_PROCESS) \ template \ class OP_NAME { \ public: \ using InterType = typename AggregateType::type; \ \ private: \ SD_HOST_DEVICE SD_INLINE static InterType op_logic(X d1, Z* params) { \ OPERATION \ } \ \ SD_HOST_DEVICE SD_INLINE static InterType merge_logic(InterType old, InterType opOutput, Z* extraParams) { \ return MERGE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static InterType update_logic(InterType old, InterType opOutput, Z* extraParams) { \ return UPDATE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static Z postProcess_logic(InterType reduction, sd::LongType n, Z* extraParams) { \ return static_cast(POST_PROCESS); \ } \ \ public: \ static const bool requiresSpecialAccumulation = false; \ \ /* Primary execSpecial signature - matches what reduce_long.hpp expects */ \ static void execSpecial(const X *x, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfoBuffer, sd::LongType *dimension, sd::LongType dimensionLength, \ const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {} \ \ SD_INLINE SD_DEVICE static void execSpecialCuda( \ const X *dx, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, \ Z *reductionBuffer, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets) {} \ \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ static SD_HOST_DEVICE X startingValue(const X* input) { return STARTING_VAL; } \ \ /* Core functions - these use Z* parameters */ \ static SD_HOST_DEVICE InterType merge(InterType old, InterType opOutput, Z* extraParams) { \ return merge_logic(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE InterType update(InterType old, InterType opOutput, Z* extraParams) { \ return update_logic(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE InterType op(X d1, Z* extraParams) { \ return op_logic(d1, extraParams); \ } \ \ static SD_HOST_DEVICE Z postProcess(InterType reduction, sd::LongType n, Z* extraParams) { \ return postProcess_logic(reduction, n, extraParams); \ } \ \ /* Template overloads for different parameter types - only when types differ */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ op(X d1, E* extraParams) { \ return op_logic(d1, static_cast(nullptr)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ update(InterType old, InterType opOutput, E* extraParams) { \ return update_logic(old, opOutput, static_cast(nullptr)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ merge(InterType old, InterType opOutput, E* extraParams) { \ return merge_logic(old, opOutput, static_cast(nullptr)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, Z> \ postProcess(InterType reduction, sd::LongType n, E* extraParams) { \ return postProcess_logic(reduction, n, static_cast(nullptr)); \ } \ \ /* FIXED: Conditional overload for sd::LongType* - only when Z is NOT sd::LongType */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && std::is_void_v, Z> \ postProcess(InterType reduction, sd::LongType n, sd::LongType* extraParams) { \ return postProcess_logic(reduction, n, static_cast(nullptr)); \ } \ \ /* Type conversion overloads for InterType mismatches */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, Z> \ postProcess(T reduction, sd::LongType n, Z* extraParams) { \ return postProcess(static_cast(reduction), n, extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ update(T old, InterType opOutput, Z* extraParams) { \ return update(static_cast(old), opOutput, extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ update(InterType old, T opOutput, Z* extraParams) { \ return update(old, static_cast(opOutput), extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, InterType> \ update(T old, U opOutput, Z* extraParams) { \ return update(static_cast(old), static_cast(opOutput), extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ merge(T old, InterType opOutput, Z* extraParams) { \ return merge(static_cast(old), opOutput, extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ merge(InterType old, T opOutput, Z* extraParams) { \ return merge(old, static_cast(opOutput), extraParams); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, InterType> \ merge(T old, U opOutput, Z* extraParams) { \ return merge(static_cast(old), static_cast(opOutput), extraParams); \ } \ \ /* Overloads for X* parameter types when X != Z */ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ op(X d1, X* extraParams) { \ return op(d1, reinterpret_cast(extraParams)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, Z> \ postProcess(InterType reduction, sd::LongType n, X* extraParams) { \ return postProcess(reduction, n, reinterpret_cast(extraParams)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t && !std::is_same_v, Z> \ postProcess(T reduction, sd::LongType n, X* extraParams) { \ return postProcess(static_cast(reduction), n, reinterpret_cast(extraParams)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ update(InterType old, InterType opOutput, X* extraParams) { \ return update(old, opOutput, reinterpret_cast(extraParams)); \ } \ \ template \ static SD_HOST_DEVICE typename std::enable_if_t, InterType> \ merge(InterType old, InterType opOutput, X* extraParams) { \ return merge(old, opOutput, reinterpret_cast(extraParams)); \ } \ }; #define DECLARE_REDUCE_LONG_OP_WITH_TYPE_CONVERSION(OP_NAME, OPERATION, REDUCE_TYPE_VAL, STARTING_VAL, MERGE_OP, UPDATE_OP, POST_PROCESS) \ template \ class OP_NAME { \ private: \ SD_HOST_DEVICE SD_INLINE static Z op_logic(X d1, X* params) { \ OPERATION \ } \ \ /* Template overload for sd::LongType* extraParams */ \ SD_HOST_DEVICE SD_INLINE static Z op_logic_converted(X d1, sd::LongType* params) { \ if (params == nullptr) return op_logic(d1, static_cast(nullptr)); \ X convertedParams[3] = {0, 0, 0}; \ for (int i = 0; i < 3; i++) { \ convertedParams[i] = static_cast(params[i]); \ } \ return op_logic(d1, convertedParams); \ } \ \ SD_HOST_DEVICE SD_INLINE static Z merge_logic(Z old, Z opOutput, X* extraParams) { \ return MERGE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static Z update_logic(Z old, Z opOutput, X* extraParams) { \ return UPDATE_OP; \ } \ \ SD_HOST_DEVICE SD_INLINE static Z postProcess_logic(Z reduction, sd::LongType n, X* extraParams) { \ return POST_PROCESS; \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd(TX d1, TX* params) { \ return op_logic(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd(TX d1, TX* params) { \ return op_logic(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_safe op_simd_converted(TX d1, sd::LongType* params) { \ return op_logic_converted(d1, params); \ } \ \ template \ SD_HOST_DEVICE SD_INLINE static enable_if_simd_unsafe op_simd_converted(TX d1, sd::LongType* params) { \ return op_logic_converted(d1, params); \ } \ \ public: \ /* Standard special execution declarations */ \ static const bool requiresSpecialAccumulation = false; \ \ static void execSpecial(const X *x, const sd::LongType *xShapeInfo, Z *extraParams, Z *result, \ const sd::LongType *resultShapeInfoBuffer, sd::LongType *dimension, sd::LongType dimensionLength, \ const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {} \ \ /* Template overload for sd::LongType* extraParams */ \ static void execSpecial(const X *x, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfoBuffer, sd::LongType *dimension, sd::LongType dimensionLength, \ const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {} \ \ /* CUDA versions */ \ SD_INLINE SD_DEVICE static void execSpecialCuda( \ const X *dx, const sd::LongType *xShapeInfo, Z *extraParams, Z *result, \ const sd::LongType *resultShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, \ Z *reductionBuffer, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets) {} \ \ SD_INLINE SD_DEVICE static void execSpecialCuda( \ const X *dx, const sd::LongType *xShapeInfo, sd::LongType *extraParams, Z *result, \ const sd::LongType *resultShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, \ Z *reductionBuffer, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets) {} \ \ const static functions::ReduceType reduceType = functions::ReduceType::REDUCE_TYPE_VAL; \ \ static SD_HOST_DEVICE X startingValue(const X* input) { return STARTING_VAL; } \ \ /* Primary op function with X* extraParams */ \ static SD_HOST_DEVICE Z op(X d1, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic(d1, extraParams); \ else \ return op_simd(d1, extraParams); \ } \ \ /* Overload for sd::LongType* extraParams */ \ static SD_HOST_DEVICE Z op(X d1, sd::LongType* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value || \ simdOps::is_simd_unsupported_argument_type::value) \ return op_logic_converted(d1, extraParams); \ else \ return op_simd_converted(d1, extraParams); \ } \ \ static SD_HOST_DEVICE Z merge(Z old, Z opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return merge_logic(old, opOutput, extraParams); \ else \ return merge_logic(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE Z update(Z old, Z opOutput, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return update_logic(old, opOutput, extraParams); \ else \ return update_logic(old, opOutput, extraParams); \ } \ \ static SD_HOST_DEVICE Z postProcess(Z reduction, sd::LongType n, X* extraParams) { \ if constexpr (simdOps::is_simd_unsupported_return_type::value) \ return postProcess_logic(reduction, n, extraParams); \ else \ return postProcess_logic(reduction, n, extraParams); \ } \ }; #define DECLARE_COUNT_REDUCE_OP(OP_NAME, CONDITION, STARTING_VAL) \ DECLARE_REDUCE_LONG_OP_WITH_TYPE_CONVERSION(OP_NAME, \ return (CONDITION) ? static_cast(1) : static_cast(0);, \ SUM, STARTING_VAL, \ old + opOutput, \ old + opOutput, \ reduction \ ) #define DECLARE_MATCH_CONDITION_REDUCE_OP(OP_NAME, MATCH_LOGIC) \ DECLARE_REDUCE_LONG_OP_WITH_TYPE_CONVERSION(OP_NAME, \ MATCH_LOGIC, \ SUM, static_cast(0), \ old + opOutput, \ old + opOutput, \ reduction \ ) } // namespace simdOps #endif // OP_MACROS_REDUCE_H_