ANSHUMAN87 commented on a change in pull request #6066: URL: https://github.com/apache/incubator-tvm/pull/6066#discussion_r459511992
########## 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); Review comment: Thanks for bringing this up! In fact i also had the same idea initially, and it was working perfectly fine. But during CI failures i discovered, there are some dependency on the positioning of these If statements with global scope variables. So i have to put this logic to avoid hoisting for any such cases. But as i mentioned in the comment as well, if you have any specific case to enable hoisting, we can add it, provided it does not violate other Pass logic. Please let me know your thought on this. ---------------------------------------------------------------- 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]
