roastduck commented on a change in pull request #6066: URL: https://github.com/apache/incubator-tvm/pull/6066#discussion_r459838669
########## 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 Object*, const Object*>; +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)) { Review comment: In both cases above, `target_for_pos` ends up to be the inner-most loop, i.e., equals to `n`. I think we only have to check whether `target_for_pos < n`, instead of running `IsParentForLoop`. Am I right? ---------------------------------------------------------------- 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: us...@infra.apache.org
