ANSHUMAN87 commented on a change in pull request #6066: URL: https://github.com/apache/incubator-tvm/pull/6066#discussion_r459859824
########## 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: I found that the Passes working on these Attr nodes are distributed. Also the if nodes are likely optimized out in those cases. But in reply to your concern, my current implementation does not stop hoisting for all Attr variables, it considers the required ones. It consider the if condition which has mixed case like `if(i + global_var == 3)` . So i think your actual concern is addressed. But if you have any specific case, which current logic does not handle, please let me know. I can check on it. Thanks! ---------------------------------------------------------------- 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]
