merrymercy commented on a change in pull request #5962: URL: https://github.com/apache/incubator-tvm/pull/5962#discussion_r449743363
########## File path: src/ansor/search_task.h ########## @@ -0,0 +1,153 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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. + */ + +/*! + * \file ansor/search_task.h + * \brief Meta information and hardware parameters for a search task. + */ + +#ifndef TVM_ANSOR_SEARCH_TASK_H_ +#define TVM_ANSOR_SEARCH_TASK_H_ + +#include <tvm/target/target.h> + +#include "compute_dag.h" + +namespace tvm { +namespace ansor { + +class HardwareParams; + +/*! \brief The parameters of target hardware used to guide the search process of SearchPolicy. */ +class HardwareParamsNode : public Object { + public: + /*! \brief The number of cores. */ + int num_cores; + /*! \brief The width of vector units in bytes. */ + int vector_unit_bytes; + /*! \brief The size of cache line in bytes. */ + int cache_line_bytes; + + // Some GPU related limitations + // Get from TVM device api + + /*! \brief The max shared memory per block. */ + int max_shared_memory_per_block{INT32_MAX}; + /*! \brief The max register memory per block. */ + int max_registers_per_block{INT32_MAX}; + /*! \brief The max threads per block. */ + int max_threads_per_block{INT32_MAX}; + /*! \brief The max vthread extent. */ + int max_vthread_extent{INT32_MAX}; + /*! \brief The thread numbers of a warp. */ + int warp_size{INT32_MAX}; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("num_cores", &num_cores); + v->Visit("vector_unit_bytes", &vector_unit_bytes); + v->Visit("cache_line_bytes", &cache_line_bytes); + v->Visit("max_shared_memory_per_block", &max_shared_memory_per_block); + v->Visit("max_registers_per_block", &max_registers_per_block); + v->Visit("max_threads_per_block", &max_threads_per_block); + v->Visit("max_vthread_extent", &max_vthread_extent); + v->Visit("warp_size", &warp_size); + } + + /*! + * \brief Get the default hardware params. + * \param target A `tvm.target`. + * \param target_host A `tvm.target` for host device. + * \return A HardwareParams object. + */ + static HardwareParams GetDefaultHardwareParams(const Target& target, const Target& target_host); + + static constexpr const char* _type_key = "ansor.HardwareParams"; + TVM_DECLARE_FINAL_OBJECT_INFO(HardwareParamsNode, Object); +}; + +/*! + * \brief Managed reference to HardwareParamsNode. + * \sa HardwareParamsNode + */ +class HardwareParams : public ObjectRef { + public: + /*! + * \brief The constructor. + * \param num_cores The number of cores. + * \param vector_unit_bytes The width of vector units in bytes. + * \param cache_line_bytes The size of cache line in bytes. + */ + HardwareParams(int num_cores, int vector_unit_bytes, int cache_line_bytes); + + TVM_DEFINE_OBJECT_REF_METHODS(HardwareParams, ObjectRef, HardwareParamsNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(HardwareParamsNode); +}; + +/*! + * \brief The computation information and hardware parameters for a specific schedule search task. + */ +class SearchTaskNode : public Object { + public: + /*! \brief The ComputeDAG for target compute declaration. */ Review comment: replace all "target compute declaration" with "input compute declaration" in all files ########## File path: src/ansor/search_policy/search_policy.h ########## @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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. + */ + +/*! + * \file ansor/search_policy/search_policy.h + * \brief The base class of search policies, including the abstract definition of search policy and + * other supporting data structures. + * + * The basic schedule search process for Ansor is design to be: + * `Program sampling` -> `Performance Tuning`. + * + * In `Program sampling`, we use some predefined precise or heuristic rules to generate several + * initial schedules. Based on these initial starting points, we perform `Performance Tuning` which + * uses cost model based evolutionary search to select schedules with the best performance. + * + * Candidate schedules are measured against the specific hardware target. + * + * \note Adding a new search policy. + * In design, there's no need for users to implement their own search policy, our formal search + * policy(will be brought later) should be enough to cover most use cases. Meanwhile, a custom rule + * mechanism will be provided to enable user-defined template search to serve the same functionality + * as the current AutoTVM template. + * + * This guide is to help understand it better and incase some advanced users have special + * requirements. + * 1. The only funcion that must be implemented is Search(), the design principe for it is to be + * the entry of starting a schedule search process and returns the best schedule get. + * 2. Information about the compute declaration of ops/subgraphs can be acquired from SearchTask. + * This structure also contains some information about the target device. (e.g. knowing the weight + * of the device vector unit, we can limit the max vectorize size during schedule generating) + * 3. SearchCallback provides more flexibility to do extra affairs during the search process. + * 4. ProgramMeasurer provides a simple but useful api to help check the performance of states get + * during the search process. + */ + +#ifndef TVM_ANSOR_SEARCH_POLICY_SEARCH_POLICY_H_ +#define TVM_ANSOR_SEARCH_POLICY_SEARCH_POLICY_H_ + +#include <tvm/node/node.h> + +#include <unordered_set> +#include <vector> + +#include "../search_task.h" + +namespace tvm { +namespace ansor { + +class ProgramMeasurer; +class SearchPolicyNode; + +/*! + * \brief Callback function to be called by the search process. + * This interface allows to do extra initializations before schedule search or extra + * check during/after the schedule search. + */ +class SearchCallbackNode : public Object { + public: + /*! + * \brief Run the registered callback function. + * \param policy A pointer to a SearchPolicyNode. + */ + virtual void Callback(SearchPolicyNode* policy) = 0; + + static constexpr const char* _type_key = "ansor.SearchCallback"; + TVM_DECLARE_BASE_OBJECT_INFO(SearchCallbackNode, Object); +}; + +/*! + * \brief Managed reference to SearchCallbackNode. + * \sa SearchCallbackNode + */ +class SearchCallback : public ObjectRef { + public: + TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(SearchCallback, ObjectRef, SearchCallbackNode); +}; + +/*! + * \brief The base class of search policies. + */ +class SearchPolicyNode : public Object { + public: + /*! \brief The current search task. */ + SearchTask cur_task; + /*! + * \brief Verbose level to control the screen output during schedule search. + * 0 for silent, 1 to output information. + */ + int verbose; + + void VisitAttrs(AttrVisitor* v) { + v->Visit("cur_task", &cur_task); + v->Visit("verbose", &verbose); + } + + /*! + * \brief Do schedule search for a task. Takes the SearchTask as input and returns the best state + * get during the search process. + * \param task The target search task. Review comment: replace `target search task` with `search task` in all files ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
