/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/epilogue/chainer/epilogue_chainer.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/epilogue/chainer/epilogue_chainer.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/epilogue/chainer/epilogue_chainer.hpp Source File
Go to the documentation of this file.
 // Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 // SPDX-License-Identifier: MIT
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common/tensor_layout.hpp"
  
 namespace ck_tile {
  
 template <typename Scheduler>
 struct EpilogueChainer
 {
     using Problem = typename Scheduler::ProblemType;
     using BaseOp  = typename Scheduler::BaseOp;
  
     using ODataType                       = typename BaseOp::ODataType;
     using DsDataType                      = typename BaseOp::DsDataType;
     using DsLayout                        = typename BaseOp::DsLayout;
     using AccDataType                     = typename BaseOp::AccDataType;
     static constexpr auto MemoryOperation = BaseOp::MemoryOperation;
  
     CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return BaseOp::GetSmemSize(); }
  
     CK_TILE_HOST_DEVICE static constexpr index_t GetVectorSizeC()
     {
         return BaseOp::GetVectorSizeC();
     }
  
     template <index_t I>
     CK_TILE_HOST_DEVICE static constexpr index_t GetVectorSizeD(number<I> idx)
     {
         return BaseOp::GetVectorSizeD(idx);
     }
  
     CK_TILE_DEVICE static constexpr auto MakeLdsDistributionEncode()
     {
         return BaseOp::MakeLdsDistributionEncode();
     }
  
     template <typename OutWindow, typename AccTile, typename AuxWindows, typename... Args>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Args&&... args) const
     {
         // The context serves as a shared workspace that maintains intermediate results
         // and resources across multiple epilogue operations.
         auto context  = Scheduler::create_context(out_window, acc_tile, aux_windows, p_smem);
         auto schedule = Scheduler::make_schedule(std::forward<Args>(args)...);
         schedule(out_window, acc_tile, aux_windows, p_smem, context);
     }
 };
  
 template <typename EpilogueType, typename... Args>
 struct EpilogueNode
 {
     using Epilogue = EpilogueType;
     ck_tile::tuple<Args...> args;
  
     constexpr EpilogueNode(Args... a) : args(a...) {}
  
     template <typename OutWindow, typename AccTile, typename AuxWindows, typename Context>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Context& context) const
     {
         ck_tile::apply(
             [&](auto&&... epilogue_args) {
                 EpilogueType{}(out_window,
                                acc_tile,
                                aux_windows,
                                p_smem,
                                context,
                                std::forward<decltype(epilogue_args)>(epilogue_args)...);
             },
             args);
     }
  
     template <typename OutWindow,
               typename AccTile,
               typename AuxWindows,
               typename Context,
               index_t I>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Context& context,
                                    number<I> iAccess) const
     {
         ck_tile::apply(
             [&](auto&&... epilogue_args) {
                 EpilogueType{}(out_window,
                                acc_tile,
                                aux_windows,
                                p_smem,
                                iAccess,
                                context,
                                std::forward<decltype(epilogue_args)>(epilogue_args)...);
             },
             args);
     }
 };
  
 template <typename EpilogueType>
 struct EpilogueNode<EpilogueType>
 {
     using Epilogue = EpilogueType;
     ck_tile::tuple<> args;
  
     constexpr EpilogueNode() = default;
  
     template <typename OutWindow, typename AccTile, typename AuxWindows, typename Context>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Context& context) const
     {
         EpilogueType{}(out_window, acc_tile, aux_windows, p_smem, context);
     }
  
     template <typename OutWindow,
               typename AccTile,
               typename AuxWindows,
               typename Context,
               index_t I>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Context& context,
                                    number<I> iAccess) const
     {
         EpilogueType{}(out_window, acc_tile, aux_windows, p_smem, iAccess, context);
     }
 };
  
 template <index_t Steps, typename... EpilogueTypes>
 struct EpilogueGraph
 {
     ck_tile::tuple<EpilogueTypes...> epilogues;
  
     constexpr EpilogueGraph() = default;
     constexpr EpilogueGraph(EpilogueTypes... eps) : epilogues(eps...) {}
  
     template <typename OutWindow, typename AccTile, typename AuxWindows, typename Context>
     CK_TILE_DEVICE void operator()(OutWindow& out_window,
                                    const AccTile& acc_tile,
                                    const AuxWindows& aux_windows,
                                    void* p_smem,
                                    Context& context) const
     {
         // For each iteration, process all epilogues in order
         static_for<0, Steps, 1>{}([&](auto iAccess) {
             static_for<0, sizeof...(EpilogueTypes), 1>{}([&](auto I) {
                 epilogues.template get<I.value>()(
                     out_window, acc_tile, aux_windows, p_smem, context, iAccess);
             });
         });
     }
 };
  
 template <typename EpilogueType, typename... Args>
 constexpr auto make_node(Args... args)
 {
     return EpilogueNode<EpilogueType, Args...>{args...};
 }
  
 template <index_t Steps, typename... EpilogueTypes>
 constexpr auto make_graph(EpilogueTypes... epilogues)
 {
     return EpilogueGraph<Steps, EpilogueTypes...>{epilogues...};
 }
  
 } // namespace ck_tile