text: /**
   * \file src/opr/include/megbrain/opr/search_policy/algo_chooser.h
   * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
   *
   * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
   * implied.
   */
  
  #pragma once
  
  #include "megbrain/opr/search_policy/profiler.h"
  #include "megbrain/opr/dnn/convolution.h"
  
  template <class MegDNNOpr>
  struct MegDNNOpr2MGBOpr;
  
  #define cb(_Opr)                            \
      template <>                             \
      struct MegDNNOpr2MGBOpr<megdnn::_Opr> { \
          using MGBOpr = mgb::opr::_Opr;      \
      };
  
  MGB_FOREACH_FASTRUN_OPR(cb)
  
  #undef cb
  
  namespace mgb {
  namespace opr {
  
  /* =================== AlgoChooser =================== */
  /*!
   * \brief choose algorithm according to ExecutionPolicy
   *
   * This class only provides static methods, and the entry point is
   * AlgoChooser::setup_algo. When profiling is needed, it would first try to
   * retrive profiling stats from cache, and run TimedProfiler when necessary
   *
   * \tparam Opr megdnn operator impl
   */
  template <typename Opr>
  class AlgoChooser {
      static constexpr int arity_in = OprArityTrait<Opr>::arity_in;
      static constexpr int arity_out = OprArityTrait<Opr>::arity_out;
      static constexpr int arity = OprArityTrait<Opr>::arity;
  
      using ImplAlgo = typename Opr::AlgorithmInfo;
      using MGBOpr = typename MegDNNOpr2MGBOpr<Opr>::MGBOpr;
      using TensorLayoutArray = std::array<TensorLayout, arity>;
  
      class ExeContext {
          const TensorLayoutArray& m_layouts;
          Opr* m_megdnn_opr;
          const MGBOpr* m_mgb_opr;
          bool m_allow_weight_preprocess;
  
      public:
          ExeContext(const TensorLayoutArray& layouts, Opr* megdnn_opr,
                     const MGBOpr* mgb_opr, bool allow_weight_preprocess)
                  : m_layouts{layouts},
                    m_megdnn_opr{megdnn_opr},
                    m_mgb_opr{mgb_opr},
                    m_allow_weight_preprocess{allow_weight_preprocess} {
              mgb_assert(m_layouts.size() == layouts.size());
              static_assert(
                      std::tuple_size<TensorLayoutArray>::value == 3 ||
                              std::tuple_size<TensorLayoutArray>::value == 5 ||
                              std::tuple_size<TensorLayoutArray>::value == 8,
                      "Convolution AlgoChooser assumes arity = 3 , 5 or 8 (for "
                      "deformable conv)");
          }
  
          Opr* megdnn_opr() const { return m_megdnn_opr; }
  
          const MGBOpr* mgb_opr() const { return m_mgb_opr; }
  
          const TensorLayout& inp_layout(size_t idx) const {
              return m_layouts[idx];
          }
  
          const TensorLayoutArray& layouts() const { return m_layouts; }
  
          ImplAlgo choose_by_heuristic(bool reproducible = false) const;
  
          //! get all candidate algos, and the one choose_by_heuristic() is
          //! put first
          std::vector<ImplAlgo> get_all_candidates() const;
  
          //! get candidate algos with workspace limit.
          std::vector<ImplAlgo> get_all_candidates_with_workspace_limit() const;
  
          //! get workspace size required for specific algo
          size_t get_workspace_size_bytes(ImplAlgo algo) const;
  
          /*!
           * \brief profile a single algorithm
           *
           * This is actually a wrapper that constructs param and call
           * TimedProfiler<Opr>::profile for the actual profiling
           *
           * \param[in,out] timeout set the timeout, and return the actual
           *      timeout used during profiling
           */
          Maybe<AlgoChooserProfileCache::ResultEntry> profile_single_algo(
                  ImplAlgo algo, double& timeout) const;
  
      private:
          Maybe<PreprocessFilter<Opr>> construct_fake_preprocess_filter() const;
      };
  
      //! entrance for getting algorithm according to execution strategy
      static ImplAlgo get_algo(ExeContext& ctx);
  
      //! get all profile result, either by retrieving cache or profiling
      static AlgoChooserProfileCache::Result get_profile_result(
              ExeContext& ctx, bool enable_update);
  
      static ImplAlgo choose_by_profile(ExeContext& ctx,
                                        bool require_reproducible,
                                        bool enable_update = true);
  
  public:
      /*!
       * \brief setup algorithm and return workspace size
       */
      static size_t setup_algo(const TensorLayoutArray& layouts, Opr* megdnn_opr,
                               const MGBOpr* mgb_opr,
                               bool allow_weight_preprocess = false);
  };
  
  }  // namespace opr
  }  // namespace mgb
  
  // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}
  