You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/07/24 14:50:10 UTC
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull request #6066: [TIR][Transform] HoistIfThenElse added
MarisaKirisame commented on a change in pull request #6066:
URL: https://github.com/apache/incubator-tvm/pull/6066#discussion_r460100741
##########
File path: src/tir/transforms/hoist_if_then_else.cc
##########
@@ -0,0 +1,376 @@
+/*
+ * 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 hoist_if_then_else.cc
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt_functor.h>
+#include <tvm/tir/transform.h>
+
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "../../arith/interval_set.h"
+#include "../../runtime/thread_storage_scope.h"
+#include "ir_util.h"
+
+namespace tvm {
+namespace tir {
+
+using VarForMap = std::unordered_map<const VarNode*, const ForNode*>;
+using HoistForIfTuple = std::tuple<bool, const ForNode*, const IfThenElseNode*>;
+
+/*
+ * This pass tries to hoist IfThenElse stmt out of For loop if condition is loop invariant.
+ * For example, given the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * We first detect all IfThenElse stmt and find the corresponding loop invariant For stmt.
+ * Then we hoist IfThenElse stmt by one For stmt each step:
+ *
+ * Step 1:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Step 2:
+ * for (i = 0; i < 3; i++)
+ * if (likely(i*2 < 4))
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * In this pass, we only continue detecting possible hoisting chance when visiting For,
+ * IfThenElse or AttrStmt Node. For example, for the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Only the For with k variable will be considered and the resulting stmt would be:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * This pass doesn't do hoisting for consecutive IfThenElse stmt. The following
+ * block won't be optimized:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ * if (likely(j > 2))
+ * A[i+j+k] = B[i+j+k]
+ *
+ */
+
+// Select potential candidate IRs that can be hoisted.
+class HoistCandidateSelector final : public StmtExprVisitor {
+ public:
+ HoistCandidateSelector() { InitRecorder(); }
+
+ void VisitStmt_(const ForNode* op) final {
+ // Check if it is first for loop, then start the recorder
+ if (!RecordingComplete()) {
+ StartOrAddRecord(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveRecord(op);
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const SeqStmtNode* op) final {
+ // If SeqStmt is encountered in the middle of recording
+ // then need to purge all, as it can not be hoisted
+ if (IsRecordingOn()) {
+ ResetRecorder();
+ }
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) final {
+ // Maintain list of all vars in AttrStmt
+ // To stop hoisting if any of the block variables are used.
+ //
+ // NOTE: If in future
+ // hoisting is required for any specific case,
+ // then add exception to only those case
+ // rather than allowing for all.
+ UpdateAttrVarList(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveAttrVarList(op);
+ }
+
+ void VisitStmt_(const IfThenElseNode* op) final {
+ if (IsRecordingOn()) {
+ is_if_cond = true;
+ StmtExprVisitor::VisitExpr(op->condition);
+ is_if_cond = false;
+
+ if (CheckValidIf()) {
+ // Check corresponding for loop
+ bool match_found = false;
+ size_t match_for_loop_pos = 0;
+ for (auto var : if_var_list_) {
+ for (size_t i = 0; i < ordered_for_list_.size() - 1; ++i) {
+ if (ordered_for_list_[i] == var_for_map_[var]) {
+ if (match_for_loop_pos < i) {
+ match_for_loop_pos = i;
+ }
+ match_found = true;
+ break;
+ }
+ }
+ }
+ // If none of the for loop has the matching loop variable as if condition,
+ // then the if node need to be hoisted on top of all, provided no parent loop exists.
+ int target_for_pos = match_found ? match_for_loop_pos + 1 : 0;
+
+ // Check if target for loop is not the parent of current if node
+ if (!IsParentForLoop(target_for_pos)) {
+ StopAndAddRecord(ordered_for_list_[target_for_pos], op);
+ }
+ }
+ if_var_list_.clear();
+ StmtExprVisitor::VisitStmt_(op);
+ StopRecording();
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ if (is_if_cond) {
+ if_var_list_.emplace_back(op);
+ }
+ }
+
+ HoistForIfTuple hoist_for_if_recorder;
+
+ void ResetRecorder() {
+ if (is_recorder_on) {
+ CHECK_GT(ordered_for_list_.size(), 0);
+ is_recorder_on = false;
+ }
+ ordered_for_list_.clear();
+ var_for_map_.clear();
+ hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr);
+ }
+
+ bool RecordingComplete() {
+ if (std::get<0>(hoist_for_if_recorder)) return true;
+ return false;
+ }
+
+ const ForNode* GetTargetForNode() { return std::get<1>(hoist_for_if_recorder); }
+
+ const IfThenElseNode* GetTargetIfNode() { return std::get<2>(hoist_for_if_recorder); }
+
+ private:
+ bool CheckValidIf() {
+ // If no if var list is present, then all the condition vars are possibly from AttrStmt, so stop
+ // hoisting
+ if (if_var_list_.size() == 0) {
+ return false;
+ }
+ if (CheckAttrVar()) {
+ return false;
+ }
+ return true;
+ }
+
+ bool IsParentForLoop(int loop_pos) {
+ // Check if the loop position is higher than the parent loop position
+ for (auto var : if_var_list_) {
+ if (GetParentLoopPos(var_for_map_[var]) >= loop_pos) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ int GetParentLoopPos(const Object* node) {
+ for (size_t i = 0; i < ordered_for_list_.size(); ++i) {
+ if (ordered_for_list_[i] == node) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ void InitRecorder() { hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr); }
+
+ void StopRecording() {
+ if (is_recorder_on) is_recorder_on = false;
Review comment:
```suggestion
is_recorder_on = false;
```
##########
File path: src/tir/transforms/hoist_if_then_else.cc
##########
@@ -0,0 +1,376 @@
+/*
+ * 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 hoist_if_then_else.cc
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt_functor.h>
+#include <tvm/tir/transform.h>
+
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "../../arith/interval_set.h"
+#include "../../runtime/thread_storage_scope.h"
+#include "ir_util.h"
+
+namespace tvm {
+namespace tir {
+
+using VarForMap = std::unordered_map<const VarNode*, const ForNode*>;
+using HoistForIfTuple = std::tuple<bool, const ForNode*, const IfThenElseNode*>;
+
+/*
+ * This pass tries to hoist IfThenElse stmt out of For loop if condition is loop invariant.
+ * For example, given the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * We first detect all IfThenElse stmt and find the corresponding loop invariant For stmt.
+ * Then we hoist IfThenElse stmt by one For stmt each step:
+ *
+ * Step 1:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Step 2:
+ * for (i = 0; i < 3; i++)
+ * if (likely(i*2 < 4))
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * In this pass, we only continue detecting possible hoisting chance when visiting For,
+ * IfThenElse or AttrStmt Node. For example, for the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Only the For with k variable will be considered and the resulting stmt would be:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * This pass doesn't do hoisting for consecutive IfThenElse stmt. The following
+ * block won't be optimized:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ * if (likely(j > 2))
+ * A[i+j+k] = B[i+j+k]
+ *
+ */
+
+// Select potential candidate IRs that can be hoisted.
+class HoistCandidateSelector final : public StmtExprVisitor {
+ public:
+ HoistCandidateSelector() { InitRecorder(); }
+
+ void VisitStmt_(const ForNode* op) final {
+ // Check if it is first for loop, then start the recorder
+ if (!RecordingComplete()) {
+ StartOrAddRecord(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveRecord(op);
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const SeqStmtNode* op) final {
+ // If SeqStmt is encountered in the middle of recording
+ // then need to purge all, as it can not be hoisted
+ if (IsRecordingOn()) {
+ ResetRecorder();
+ }
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) final {
+ // Maintain list of all vars in AttrStmt
+ // To stop hoisting if any of the block variables are used.
+ //
+ // NOTE: If in future
+ // hoisting is required for any specific case,
+ // then add exception to only those case
+ // rather than allowing for all.
+ UpdateAttrVarList(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveAttrVarList(op);
+ }
+
+ void VisitStmt_(const IfThenElseNode* op) final {
+ if (IsRecordingOn()) {
+ is_if_cond = true;
+ StmtExprVisitor::VisitExpr(op->condition);
+ is_if_cond = false;
+
+ if (CheckValidIf()) {
+ // Check corresponding for loop
+ bool match_found = false;
+ size_t match_for_loop_pos = 0;
+ for (auto var : if_var_list_) {
+ for (size_t i = 0; i < ordered_for_list_.size() - 1; ++i) {
+ if (ordered_for_list_[i] == var_for_map_[var]) {
+ if (match_for_loop_pos < i) {
+ match_for_loop_pos = i;
+ }
+ match_found = true;
+ break;
+ }
+ }
+ }
+ // If none of the for loop has the matching loop variable as if condition,
+ // then the if node need to be hoisted on top of all, provided no parent loop exists.
+ int target_for_pos = match_found ? match_for_loop_pos + 1 : 0;
+
+ // Check if target for loop is not the parent of current if node
+ if (!IsParentForLoop(target_for_pos)) {
+ StopAndAddRecord(ordered_for_list_[target_for_pos], op);
+ }
+ }
+ if_var_list_.clear();
+ StmtExprVisitor::VisitStmt_(op);
+ StopRecording();
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ if (is_if_cond) {
+ if_var_list_.emplace_back(op);
+ }
+ }
+
+ HoistForIfTuple hoist_for_if_recorder;
+
+ void ResetRecorder() {
+ if (is_recorder_on) {
+ CHECK_GT(ordered_for_list_.size(), 0);
+ is_recorder_on = false;
+ }
+ ordered_for_list_.clear();
+ var_for_map_.clear();
+ hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr);
+ }
+
+ bool RecordingComplete() {
+ if (std::get<0>(hoist_for_if_recorder)) return true;
+ return false;
+ }
+
+ const ForNode* GetTargetForNode() { return std::get<1>(hoist_for_if_recorder); }
+
+ const IfThenElseNode* GetTargetIfNode() { return std::get<2>(hoist_for_if_recorder); }
+
+ private:
+ bool CheckValidIf() {
+ // If no if var list is present, then all the condition vars are possibly from AttrStmt, so stop
+ // hoisting
+ if (if_var_list_.size() == 0) {
+ return false;
+ }
+ if (CheckAttrVar()) {
+ return false;
+ }
+ return true;
+ }
+
+ bool IsParentForLoop(int loop_pos) {
+ // Check if the loop position is higher than the parent loop position
+ for (auto var : if_var_list_) {
+ if (GetParentLoopPos(var_for_map_[var]) >= loop_pos) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ int GetParentLoopPos(const Object* node) {
+ for (size_t i = 0; i < ordered_for_list_.size(); ++i) {
+ if (ordered_for_list_[i] == node) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ void InitRecorder() { hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr); }
+
+ void StopRecording() {
+ if (is_recorder_on) is_recorder_on = false;
+ }
+
+ bool IsRecordingOn() { return is_recorder_on; }
+
+ void StartOrAddRecord(const ForNode* op) {
+ if (!is_recorder_on) is_recorder_on = true;
+ if (!var_for_map_.count(op->loop_var.get())) {
+ var_for_map_.insert({op->loop_var.get(), op});
+ }
+ ordered_for_list_.emplace_back(op);
+ }
+
+ void RemoveRecord(const ForNode* op) {
+ StopRecording();
+ var_for_map_.erase(op->loop_var.get());
+ if (ordered_for_list_.size() > 0) ordered_for_list_.pop_back();
Review comment:
why? imo this should be a CHECK() instead.
##########
File path: src/tir/transforms/hoist_if_then_else.cc
##########
@@ -0,0 +1,376 @@
+/*
+ * 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 hoist_if_then_else.cc
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt_functor.h>
+#include <tvm/tir/transform.h>
+
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "../../arith/interval_set.h"
+#include "../../runtime/thread_storage_scope.h"
+#include "ir_util.h"
+
+namespace tvm {
+namespace tir {
+
+using VarForMap = std::unordered_map<const VarNode*, const ForNode*>;
+using HoistForIfTuple = std::tuple<bool, const ForNode*, const IfThenElseNode*>;
+
+/*
+ * This pass tries to hoist IfThenElse stmt out of For loop if condition is loop invariant.
+ * For example, given the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * We first detect all IfThenElse stmt and find the corresponding loop invariant For stmt.
+ * Then we hoist IfThenElse stmt by one For stmt each step:
+ *
+ * Step 1:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Step 2:
+ * for (i = 0; i < 3; i++)
+ * if (likely(i*2 < 4))
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * In this pass, we only continue detecting possible hoisting chance when visiting For,
+ * IfThenElse or AttrStmt Node. For example, for the following block:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * Only the For with k variable will be considered and the resulting stmt would be:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * A[i + j] = A[i + j] - 1
+ * if (likely(i*2 < 4))
+ * for (k = 0; k < 5; k++)
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ *
+ * This pass doesn't do hoisting for consecutive IfThenElse stmt. The following
+ * block won't be optimized:
+ * for (i = 0; i < 3; i++)
+ * for (j = 0; j < 4; j++)
+ * for (k = 0; k < 5; k++)
+ * if (likely(i*2 < 4))
+ * A[3*i+2j+k] = B[7*i+3j+k]
+ * if (likely(j > 2))
+ * A[i+j+k] = B[i+j+k]
+ *
+ */
+
+// Select potential candidate IRs that can be hoisted.
+class HoistCandidateSelector final : public StmtExprVisitor {
+ public:
+ HoistCandidateSelector() { InitRecorder(); }
+
+ void VisitStmt_(const ForNode* op) final {
+ // Check if it is first for loop, then start the recorder
+ if (!RecordingComplete()) {
+ StartOrAddRecord(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveRecord(op);
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const SeqStmtNode* op) final {
+ // If SeqStmt is encountered in the middle of recording
+ // then need to purge all, as it can not be hoisted
+ if (IsRecordingOn()) {
+ ResetRecorder();
+ }
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) final {
+ // Maintain list of all vars in AttrStmt
+ // To stop hoisting if any of the block variables are used.
+ //
+ // NOTE: If in future
+ // hoisting is required for any specific case,
+ // then add exception to only those case
+ // rather than allowing for all.
+ UpdateAttrVarList(op);
+ StmtExprVisitor::VisitStmt_(op);
+ RemoveAttrVarList(op);
+ }
+
+ void VisitStmt_(const IfThenElseNode* op) final {
+ if (IsRecordingOn()) {
+ is_if_cond = true;
+ StmtExprVisitor::VisitExpr(op->condition);
+ is_if_cond = false;
+
+ if (CheckValidIf()) {
+ // Check corresponding for loop
+ bool match_found = false;
+ size_t match_for_loop_pos = 0;
+ for (auto var : if_var_list_) {
+ for (size_t i = 0; i < ordered_for_list_.size() - 1; ++i) {
+ if (ordered_for_list_[i] == var_for_map_[var]) {
+ if (match_for_loop_pos < i) {
+ match_for_loop_pos = i;
+ }
+ match_found = true;
+ break;
+ }
+ }
+ }
+ // If none of the for loop has the matching loop variable as if condition,
+ // then the if node need to be hoisted on top of all, provided no parent loop exists.
+ int target_for_pos = match_found ? match_for_loop_pos + 1 : 0;
+
+ // Check if target for loop is not the parent of current if node
+ if (!IsParentForLoop(target_for_pos)) {
+ StopAndAddRecord(ordered_for_list_[target_for_pos], op);
+ }
+ }
+ if_var_list_.clear();
+ StmtExprVisitor::VisitStmt_(op);
+ StopRecording();
+ return;
+ }
+
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) final {
+ if (is_if_cond) {
+ if_var_list_.emplace_back(op);
+ }
+ }
+
+ HoistForIfTuple hoist_for_if_recorder;
+
+ void ResetRecorder() {
+ if (is_recorder_on) {
+ CHECK_GT(ordered_for_list_.size(), 0);
+ is_recorder_on = false;
+ }
+ ordered_for_list_.clear();
+ var_for_map_.clear();
+ hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr);
+ }
+
+ bool RecordingComplete() {
+ if (std::get<0>(hoist_for_if_recorder)) return true;
+ return false;
+ }
+
+ const ForNode* GetTargetForNode() { return std::get<1>(hoist_for_if_recorder); }
+
+ const IfThenElseNode* GetTargetIfNode() { return std::get<2>(hoist_for_if_recorder); }
+
+ private:
+ bool CheckValidIf() {
+ // If no if var list is present, then all the condition vars are possibly from AttrStmt, so stop
+ // hoisting
+ if (if_var_list_.size() == 0) {
+ return false;
+ }
+ if (CheckAttrVar()) {
+ return false;
+ }
+ return true;
+ }
+
+ bool IsParentForLoop(int loop_pos) {
+ // Check if the loop position is higher than the parent loop position
+ for (auto var : if_var_list_) {
+ if (GetParentLoopPos(var_for_map_[var]) >= loop_pos) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ int GetParentLoopPos(const Object* node) {
+ for (size_t i = 0; i < ordered_for_list_.size(); ++i) {
+ if (ordered_for_list_[i] == node) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ void InitRecorder() { hoist_for_if_recorder = std::make_tuple(false, nullptr, nullptr); }
+
+ void StopRecording() {
+ if (is_recorder_on) is_recorder_on = false;
+ }
+
+ bool IsRecordingOn() { return is_recorder_on; }
+
+ void StartOrAddRecord(const ForNode* op) {
+ if (!is_recorder_on) is_recorder_on = true;
Review comment:
```suggestion
is_recorder_on = true;
```
----------------------------------------------------------------
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:
users@infra.apache.org