| /* |
| * 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 auto_scheduler/search_task.h |
| * \brief Meta information and hardware parameters for a search task. |
| */ |
| |
| #ifndef TVM_AUTO_SCHEDULER_SEARCH_TASK_H_ |
| #define TVM_AUTO_SCHEDULER_SEARCH_TASK_H_ |
| |
| #include <tvm/auto_scheduler/compute_dag.h> |
| #include <tvm/runtime/ndarray.h> |
| #include <tvm/target/target.h> |
| |
| namespace tvm { |
| namespace auto_scheduler { |
| |
| class HardwareParams; |
| |
| /*! \brief The parameters of target hardware used to guide the 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; |
| |
| // GPU related parameters got from device query API |
| /*! \brief The max shared memory per block in bytes. */ |
| int max_shared_memory_per_block; |
| /*! \brief The max local memory per block in bytes. */ |
| int max_local_memory_per_block; |
| /*! \brief The max number of threads per block. */ |
| int max_threads_per_block; |
| /*! \brief The max vthread extent. */ |
| int max_vthread_extent; |
| /*! \brief The thread numbers of a warp. */ |
| int warp_size; |
| |
| 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_local_memory_per_block", &max_local_memory_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 = "auto_scheduler.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. |
| * \param max_shared_memory_per_block The max amount of shared memory per block for GPU. |
| * \param max_local_memory_per_block The max amount of local memory per block for GPU. |
| * \param max_threads_per_block The max number of threads per block for GPU. |
| * \param max_vthread_extent The max extent of vthread for GPU. |
| * \param warp_size The warp size for GPU |
| */ |
| HardwareParams(int num_cores, int vector_unit_bytes, int cache_line_bytes, |
| int max_shared_memory_per_block, int max_local_memory_per_block, |
| int max_threads_per_block, int max_vthread_extent, int warp_size); |
| |
| 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 the compute declaration. */ |
| ComputeDAG compute_dag; |
| /*! \brief The workload key for the compute declaration. */ |
| String workload_key; |
| /*! \brief The description string of this task. */ |
| String desc; |
| /*! \brief The target device of this search task. */ |
| Target target; |
| /*! \brief The target host device of this search task. */ |
| Target target_host; |
| /*! \brief Hardware parameters used in this search task. */ |
| HardwareParams hardware_params; |
| /*! \brief The layout rewrite option used for measuring programs. */ |
| LayoutRewriteOption layout_rewrite_option; |
| /*! \brief Names of some user defined input data used in program measuring. */ |
| Array<String> task_input_names; |
| |
| void VisitAttrs(tvm::AttrVisitor* v) { |
| v->Visit("compute_dag", &compute_dag); |
| v->Visit("workload_key", &workload_key); |
| v->Visit("desc", &desc); |
| v->Visit("target", &target); |
| v->Visit("target_host", &target_host); |
| v->Visit("hardware_params", &hardware_params); |
| v->Visit("layout_rewrite_option", &layout_rewrite_option); |
| v->Visit("task_input_names", &task_input_names); |
| } |
| |
| static constexpr const char* _type_key = "auto_scheduler.SearchTask"; |
| TVM_DECLARE_FINAL_OBJECT_INFO(SearchTaskNode, Object); |
| }; |
| |
| /*! |
| * \brief Managed reference to SearchTaskNode. |
| * \sa SearchTaskNode |
| */ |
| class SearchTask : public ObjectRef { |
| public: |
| /*! |
| * \brief The constructor. |
| * \param compute_dag The ComputeDAG for the compute declaration. |
| * \param workload_key The workload key for the compute declaration. |
| * \param target The target device of this search task. |
| * \param target_host The target host device of this search task. |
| * \param hardware_params Hardware parameters used in this search task. |
| * \param layout_rewrite_option The layout rewrite option used for measuring programs. |
| * \param task_input_names Names of some user defined input data used in program measuring. |
| * \param desc The description string of this task. |
| */ |
| SearchTask(ComputeDAG compute_dag, String workload_key, Target target, Target target_host, |
| Optional<HardwareParams> hardware_params, LayoutRewriteOption layout_rewrite_option, |
| Array<String> task_input_names, String desc = ""); |
| |
| TVM_DEFINE_OBJECT_REF_METHODS(SearchTask, ObjectRef, SearchTaskNode); |
| }; |
| |
| } // namespace auto_scheduler |
| } // namespace tvm |
| |
| #endif // TVM_AUTO_SCHEDULER_SEARCH_TASK_H_ |