jcf94 commented on a change in pull request #6073:
URL: https://github.com/apache/incubator-tvm/pull/6073#discussion_r456319021
##########
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:
Doc updated.
----------------------------------------------------------------
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]
