merrymercy commented on a change in pull request #5962: URL: https://github.com/apache/incubator-tvm/pull/5962#discussion_r449742369
########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,375 @@ +/* + * 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/loop_state.h + * \brief The definition of the "state" in search. + * + * Each LoopState corresponds to a specific schedule for its target ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a history of transformations used to + * construct it. + * The loop structure keeps a preview of how the schedule will finally look like after lowering the + * current state (e.g. number of iterators, the extent of each iterator, the compute_at locations + * ...). During the schedule search process, the loop structure can provide search policy with + * necessary information on how to perform further operations with the current state. + * The transform history is a sequence of TransformStep which will finally be mapped to schedule + * primitives. The steps can also be used for serialization of a state. + * + * The LoopState can be seen as a lightweight loop structure IR specifically for schedule search. + * We don't use the existing TVM IR but to extend a new structure on it is because: + * 1. We want fast incremental change to the loop structures, search policy needs to get the + * immediate loop structures update rather than after TVM lowering; + * 2. We want serializable transform history for replay, backtracking, and mutation; + * 3. We may create some macro schedule primitives that represent the combination of several TVM + * schedule primitives. + * + * When the search is complete, we will lower the state to TVM IR with TVM's schedule primitives. + * Since we share a lot of common objects during search, the transformation is implemented in copy + * on write style. All objects are immutable, which is similar to TVM IR. + */ + +#ifndef TVM_ANSOR_LOOP_STATE_H_ +#define TVM_ANSOR_LOOP_STATE_H_ + +#include <tvm/runtime/container.h> + +#include <functional> + +#include "transform_step.h" + +namespace tvm { +namespace ansor { + +using namespace tvm::tir; + +class ComputeDAG; + +/*! \brief The type of a stage. */ +enum StageType { + /*! \brief A placeholder stage. */ + kPlaceholder = 0, + /*! \brief A compute stage. */ + kCompute = 1 +}; + +/*! \brief The type of compute location. */ +enum ComputeAtType { + /*! \brief Compute at root. */ + kRoot = 0, + /*! \brief Compute inlined. */ + kInlined = 1, + /*! \brief Compute at some iterator. */ + kIter = 2, +}; + +/*! \brief The type of an iterator. */ +enum IteratorType { + /*! \brief Spatial iterator. */ + kSpace = 0, + /*! \brief Reduction iterator. */ + kReduce = 1, + /*! \brief Fused spatial and reduction iterator. */ + kMixed = 2, + /*! \brief Special iterator. (e.g. virtual root iterator) */ + kSpecial = 3 +}; + +/*! \brief The type of an iterator's annotation. */ +enum IteratorAnnotation { + /*! \brief This iterator has no annotation. */ + kNone = 0, + /*! \brief This iterator has been unrolled. */ + kUnroll = 1, + /*! \brief This iterator has been vectorized. */ + kVectorize = 2, + /*! \brief This iterator has been paralleld. */ + kParallel = 3, + /*! \brief This iterator has been bind to vthread. */ + kVThread = 4, + /*! \brief This iterator has been bind to blockIdx.x. */ + kBlockX = 5, + /*! \brief This iterator has been bind to threadIdx.x. */ + kThreadX = 6, + /*! \brief This iterator has been bind to blockIdx.y. */ + kBlockY = 7, + /*! \brief This iterator has been bind to threadIdx.y. */ + kThreadY = 8, + /*! \brief This iterator has been mapped with a tensorize intrinsic. */ + kTensorized = 9 +}; + +/*! + * \brief A for loop iterator + * Similar to tvm::IterVar in `include/tvm/tir/expr.h` + */ +class IteratorNode : public Object { + public: + /*! \brief The name of this iterator. */ + String name; + /*! \brief The target range of this iterator. */ + Range range; + /*! \brief The iterator type of this iterator. */ + IteratorType iter_type; + /*! \brief The annotation type of this iterator. */ + IteratorAnnotation annotation; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("name", &name); + v->Visit("range", &range); + } + + static constexpr const char* _type_key = "ansor.Iterator"; + TVM_DECLARE_FINAL_OBJECT_INFO(IteratorNode, Object); +}; + +/*! + * \brief Managed reference to IteratorNode. + * \sa IteratorNode + */ +class Iterator : public ObjectRef { + public: + /*! + * \brief The constructor. + * \param name The name of this iterator. + * \param range The target range of this iterator. + * \param iter_type The iterator type of this iterator. + * \param annotation The annotation type of this iterator. + */ + Iterator(String name, Range range, IteratorType iter_type, IteratorAnnotation annotation); + + TVM_DEFINE_OBJECT_REF_METHODS(Iterator, ObjectRef, IteratorNode); +}; + +/*! \brief Stage-level attributes. */ +struct StageAttributes { + /*! \brief The maximum steps for the pragma `auto_unroll_max_step`. */ + int auto_unroll_max_step; + /*! \brief The storage offset for the schedule primitive `storage_align`. */ + int storage_offset; +}; + +/*! + * \brief A op stage in the compute declaration. + * Similar to te::Stage in `include/schedule.h`. + */ +class StageNode : public Object { + public: + /*! \brief The operator of this stage */ + te::Operation op; + /*! \brief The type of this stage. */ + StageType op_type; + /*! \brief The iterators in this stage. */ + Array<Iterator> iters; + /*! \brief The compute location of this stage. */ + ComputeAtType compute_at; + /*! \brief Other stage-level attributes. */ + StageAttributes attrs; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("op", &op); + v->Visit("iters", &iters); + } + + static constexpr const char* _type_key = "ansor.Stage"; + TVM_DECLARE_FINAL_OBJECT_INFO(StageNode, Object); +}; + +/*! + * \brief Managed reference to StageNode. + * \sa StageNode + */ +class Stage : public ObjectRef { + public: + /*! + * \brief The constructor. + * \param op A `te::Operation`. + */ + explicit Stage(te::Operation op); + /*! + * \brief The constructor. + * \param op A `te::Operation`. + * \param op_type The stage type of this op. + * \param iters The iterators of this op. (copy) + * \param compute_at The compute at type of this op. + * \param attrs Other stage-level attributes. + */ + Stage(te::Operation op, StageType op_type, const Array<Iterator>& iters, ComputeAtType compute_at, + StageAttributes attrs); + /*! + * \brief The constructor. + * \param op A `te::Operation`. + * \param op_type The stage type of this op. + * \param iters The iterators of this op. (move) + * \param compute_at The compute at type of this op. + * \param attrs Other stage-level attributes. + */ + Stage(te::Operation op, StageType op_type, Array<Iterator>&& iters, ComputeAtType compute_at, + StageAttributes attrs); + + TVM_DEFINE_OBJECT_REF_METHODS(Stage, ObjectRef, StageNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(StageNode); +}; + +/*! + * \brief A state in the search process. + * It consists of the current loop structure and a history of transformations used to construct it. + * Each State corresponds to a specific schedule for its target ComputeDAG. + */ +class StateNode : public Object { + public: + /*! \brief Current stages and loop structures. */ + Array<Stage> stages; + /*! \brief History transformation steps. */ + Array<Step> transform_steps; + /*! \brief Indicate whether this state has unfilled tile sizes. */ + bool complete; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("stages", &stages); + v->Visit("transform_steps", &transform_steps); + v->Visit("complete", &complete); + } + + static constexpr const char* _type_key = "ansor.State"; + TVM_DECLARE_FINAL_OBJECT_INFO(StateNode, Object); + + private: + /*! + * \brief The up-to-date ComputeDAG of this state, used for some steps that may change the + * stage structure of the ComputeDAG, for exp. CacheReadStep/CacheWriteStep(Will be added later). + * The default value is an empty ObjectRef. (means no modification to the original DAG) + */ + ObjectRef current_compute_dag; +}; + +/*! + * \brief Managed reference to StateNode. + * \sa StateNode + */ +class State : public ObjectRef { + public: + /*! + * \brief The constructor. + * \param ops `te::Operation`s for a compute declaration. + */ + explicit State(const Array<te::Operation>& ops); + + /*! + * \brief Print the state to a human readable string. + * \param delete_trivial_loop True for skipping the trivial loops. + * (undefined or extent == 1, default set to True) + * \return The human readable state structure. + */ + String ToStr(bool delete_trivial_loop = true) const; + + /*! + * \brief General do step functions with a runtime dynamic dispatcher. + * \param dag The target ComputeDAG. + */ + void DoSteps(const ComputeDAG& dag); + + /* Step APIs for State. */ + + /*! + * \brief Schedule primitive corresponds to te.reorder. + * \param stage_id The index of the target stage. Review comment: remove 'target' from the docstrings for these three functions (reorder, split and fuse). See my comments on the python file ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,375 @@ +/* + * 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/loop_state.h + * \brief The definition of the "state" in search. + * + * Each LoopState corresponds to a specific schedule for its target ComputeDAG. Review comment: Improve the docstring by using my comments on the python file. ########## File path: src/ansor/measure.h ########## @@ -0,0 +1,430 @@ +/* + * 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/measure.h + * \brief Distributed measurement infrastructure to measure the runtime costs of tensor programs. + * MeasureInput -> BuildeResult -> MeasureResult + */ + +#ifndef TVM_ANSOR_MEASURE_H_ +#define TVM_ANSOR_MEASURE_H_ + +#include <unordered_map> +#include <utility> + +#include "loop_state.h" +#include "search_task.h" + +namespace tvm { +namespace ansor { + +class SearchPolicy; +class MeasureInput; +class MeasureResult; + +/*! \brief The error code of one measurement */ +enum MeasureErrorNO { + /*! \brief No error. */ + kNoError = 0, + /*! \brief Errors happen when apply transform steps from init state. */ + kInstantiationError = 1, + /*! \brief Errors happen when compiling code on host. (when build module) */ + kCompileHostError = 2, + /*! \brief Errors happen when compiling code on device. (when load module) */ + kCompileDeviceError = 3, + /*! \brief Errors happen when run program on device. */ + kRuntimeDeviceError = 4, + /*! \brief Answer is wrong when compared to a reference output. */ + kWrongAnswerError = 5, + /*! \brief Timeout during compilation. */ + kBuildTimeoutError = 6, + /*! \brief Timeout during run. */ + kRunTimeoutError = 7, + /*! \brief Unknown error. */ + kUnknonwError = 8, +}; + +// Inputs and results of one measurement + +/*! \brief Store the input of a measurement */ +class MeasureInputNode : public Object { + public: + /*! \brief The search task. */ + SearchTask task; + /*! \brief The program state to be measured. */ + State state; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("task", &task); + v->Visit("state", &state); + } + + /*! \brief Do deep copy. */ + MeasureInput copy() const; + + static constexpr const char* _type_key = "ansor.MeasureInput"; + TVM_DECLARE_FINAL_OBJECT_INFO(MeasureInputNode, Object); +}; + +/*! + * \brief Managed reference to MeasureInputNode. + * \sa MeasureInputNode + */ +class MeasureInput : public ObjectRef { + public: + /*! + * \brief The constructor. + * \param task The target SearchTeask. Review comment: remove all "target" before "SearchTask", "State" 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]
