yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954464119
##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
// 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;
+ int insert_position = split.first_consumer_position;
+ if (index == -1) {
+ return insert_position;
+ } else if (index == -2) {
+ if (require_all_producers_visited) {
+ if (split.last_producer_position >= 0) {
+ insert_position = split.last_producer_position + 1;
+ } else {
+ insert_position = split.first_consumer_position;
+ }
+ } else if (require_all_consumers_visited) {
+ class Finder : public StmtVisitor {
+ public:
+ void VisitStmt_(const BlockRealizeNode* realize) final {
+ const BlockNode* block = realize->block.get();
+ if (producer_blocks_.count(block)) {
+ ++this->n_producers_visited_;
+ }
+ }
+
+ std::unordered_set<const StmtNode*> producer_blocks_;
+ int n_producers_visited_ = 0;
+ };
+ // adjust the inserted position by compute at order
+ for (int i = split.first_consumer_position; i - 1 >
split.last_producer_position; --i) {
+ auto blk = GetBlock(subtrees[i]);
+ if (!blk.defined()) break;
Review Comment:
I have deleted this part of the code
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
