tqchen commented on a change in pull request #5962: URL: https://github.com/apache/incubator-tvm/pull/5962#discussion_r450436581
########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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, Review comment: Reduction ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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, Review comment: Spatial ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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 { Review comment: ComputeAtKind ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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 { Review comment: enum class IteratorKind : int ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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 { Review comment: StageKind ########## File path: python/tvm/ansor/record.py ########## @@ -0,0 +1,157 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file Review comment: record -> measure_record ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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 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 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. + * \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); + + 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 list of transformation steps used to construct + * it. + * Each State corresponds to a specific schedule for its 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; Review comment: It would also be useful to think about alternative names, how about concrete? ########## File path: src/ansor/loop_state.h ########## @@ -0,0 +1,371 @@ +/* + * 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 schedule for its ComputeDAG. + * A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to + * construct the loop structure. + * 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 manipulate the current state. + * The transform history is a sequence of `TransformStep` which will finally be mapped to TVM + * schedule primitives. The steps can also be used for the 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. The 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 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 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. + * \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); + + 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 list of transformation steps used to construct + * it. + * Each State corresponds to a specific schedule for its 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; Review comment: It would be great if we can have a description about the completeness. e.g. A complete state means that all tile sizes of the state is filled. ---------------------------------------------------------------- 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]
