junrushao1994 commented on a change in pull request #6073:
URL: https://github.com/apache/incubator-tvm/pull/6073#discussion_r456257744
##########
File path: src/auto_scheduler/loop_state.h
##########
@@ -217,6 +196,68 @@ class Stage : public ObjectRef {
TVM_DEFINE_OBJECT_REF_COW_METHOD(StageNode);
};
+/*! \brief Use stage_id to represent a stage. */
+using StageKey = int;
+/*! \brief Use stage_id and iter_id to represent a iterator. */
+using IterKey = std::pair<int, int>;
+
+/*!
+ * \brief stores the compute_at relation between stages
+ * This stores a bi-directional mapping from stages and iter:
+ * 1. Stage to its attached iterator
+ * 2. Iterator to the stage attached to it
+ * You can use AttachMapNode::stage_to_attach_iter and
AttachMapNode::iter_to_attached_stages
+ * to query the relations
+ */
+class AttachMapNode : public Object {
+ public:
+ /*! \brief A Map to store the mapping of stage to its attached iterator. */
+ std::unordered_map<StageKey, IterKey> stage_to_attach_iter;
+ /*! \brief A Map to store the mapping of iterator to the stage attached to
it. */
+ std::unordered_map<IterKey, std::vector<StageKey>> iter_to_attached_stages;
+
+ static constexpr const char* _type_key = "auto_scheduler.AttachMap";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AttachMapNode, Object);
+};
+
+/*!
+ * \brief Managed reference to AttachMapNode.
+ * \sa AttachMapNode
+ */
+class AttachMap : public ObjectRef {
+ public:
+ /*!
+ * \brief Process the stage/iterator mapping after compute at.
+ * \param stage_id The index of the stage to be compute at.
+ * \param target_stage_id The index of stage that this step will compute at
to.
+ * \param target_iter_id The index of iterator in target stage that this
step will compute at to.
+ */
+ void SetComputeAtIter(int stage_id, int target_stage_id, int target_iter_id);
+ /*!
+ * \brief This is a public wrapper of `DeleteStageEntry`. To delete the
entry of a specific stage.
+ * \param stage_id The index of the stage to be compute at.
+ */
+ void DeleteStage(int stage_id);
+ /*!
+ * \brief Update the iterator relations in AttachMap.
Review comment:
I agree with Cody...Tried to understand the doc for a while but didn't
succeed without looking at the source code. That will be great if we have
clearer doc and method name :-)
##########
File path: src/auto_scheduler/utils.h
##########
@@ -89,6 +89,15 @@ inline int GetIndex(const Array<T>& array, const T&
to_locate) {
return -1;
}
+/*! \brief Delete the item in a std::vector. */
+template <typename T>
+inline void DeleteItem(std::vector<T>* array, const T& to_delete) {
Review comment:
Maybe `FindAndDeleteItem` sounds more informative?
##########
File path: src/auto_scheduler/loop_state.cc
##########
@@ -90,12 +90,69 @@ Stage::Stage(te::Operation op, StageKind op_type, const
Array<Iterator>& iters,
data_ = std::move(node);
}
+/********** AttachMap **********/
+void AttachMap::SetComputeAtIter(int stage_id, int target_stage_id, int
target_iter_id) {
+ AttachMapNode* pnode = CopyOnWrite();
+
+ // Delete the current entry of this stage
+ DeleteStageEntry(pnode, stage_id);
+
+ // Store the new relations to map
+ IterKey iter_key(target_stage_id, target_iter_id);
+ pnode->stage_to_attach_iter[stage_id] = iter_key;
+ pnode->iter_to_attached_stages[iter_key].push_back(stage_id);
+}
+
+void AttachMap::DeleteStage(int stage_id) {
+ AttachMapNode* pnode = CopyOnWrite();
+ // Delete the original stage entry
+ DeleteStageEntry(pnode, stage_id);
+}
+
+void AttachMap::UpdateIters(const std::vector<IterKey>& old_iters,
+ const std::vector<IterKey>& new_iters) {
Review comment:
Shall we move the check to the beginning of this function before CoW?
----------------------------------------------------------------
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]
