junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r955645029
##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -159,8 +162,22 @@ int FindInsertionPoint(
// Step 3. Check if there is at least one index of the position can be
inserted into
// The valid indices are: (last_producer_position, first_consumer_position]
ICHECK(split.last_producer_position < split.first_consumer_position);
- // Step 4. Return the last valid insertion point
- return split.first_consumer_position;
+ // Step 4. Return the possible insertion point according to index
+ int insert_position;
+ if (index == -1) {
+ insert_position = split.first_consumer_position;
+ } else if (index == -2) {
+ insert_position = split.last_producer_position + 1;
+ } else if (index >= 0 && index >= split.last_producer_position + 1 &&
+ index <= split.first_consumer_position) {
+ insert_position = index;
+ } else {
+ LOG(FATAL) << "Valid index:(-1, -2, [" << split.last_producer_position + 1
<< ", "
+ << split.first_consumer_position << "]), "
+ << "current index=" << index;
+ throw;
+ }
+ return insert_position;
Review Comment:
Thanks for updating the logic here! It's much clearer now!
--
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.
To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
