Predication pass optimization for nested if statements (p4lang#2083) (p…

…4lang#2231)

* Predication pass optimization for nested ifs

* Implementation without vector with predicates for each if

* ExpressionReplacer for IR::Mux

* Output test changed

* Expression Replacer preorder Mux implemented

* Implemented ExpressionChecker

* Tests output changed

* Predication pass optimization for nested if statements (p4lang#2083)Refactored Predication pass, because of the issue p4lang#2083 in which LocalCopyProp passcreates quadratically more complex code for chained if-else statements.Predication pass is optimized by looking for same assignment statements in both thenand else branches and then converting them to the ternary expression if possible.It will not be possible to convert more than one assignment if they have dependenciesbetween each other.For the basic example with one if-else statement like this:if(e)	x = 5;else	x = 10;Code will be transformed to this:x = e ? 5 : 10;New implementation of Predication pass relies on keeping track of the transformedlive assignments which will be inserted into the right place in IR, after visitingboth if and else statements of the corresponding block.ExpressionReplacer transforms expression on the right side of the assignment intoternary operator(IR::Mux) or nested ternary operators.ReplaceChecker inspects if the assignment statement could be emplaced into thenested ternary operator.

* Removed ReplaceChecker, refactored using Context

Visitor ReplaceChecker is unnecessary because it
is covered by dependecy prints.
ExpressionReplacer now uses VisitorContext class.
Removed redundant checks.

* ordered Names erasing changed

* Removed unused includes

Co-authored-by: Your Name <you@example.com>
Co-authored-by: matijalukic <matijalukictutorijal@gmail.com>

midend/predication.cpp

@@ -13,11 +13,120 @@ 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.
 */
-
 #include "predication.h"
 #include "frontends/p4/cloner.h"
 
+
 namespace P4 {
+/// convert an expression into a string that uniqely identifies the value referenced
+/// return null cstring if not a reference to a constant thing.
+static cstring expr_name(const IR::Expression *exp) {
+    if (auto p = exp->to<IR::PathExpression>())
+        return p->path->name;
+    if (auto m = exp->to<IR::Member>()) {
+        if (auto base = expr_name(m->expr))
+            return base + "." + m->member;
+    } else if (auto a = exp->to<IR::ArrayIndex>()) {
+        if (auto k = a->right->to<IR::Constant>())
+            if (auto base = expr_name(a->left))
+                return base + "." + std::to_string(k->asInt()); }
+    return cstring();
+}
+
+bool Predication::EmptyStatementRemover::preorder(IR::BlockStatement * block) {
+    for (auto iter = block->components.begin(); iter != block->components.end();) {
+        if ((*iter)->is<IR::EmptyStatement>()) {
+            iter = block->components.erase(iter);
+        } else if ((*iter)->is<IR::BlockStatement>()) {
+            visit(*iter);
+            if ((*iter)->to<IR::BlockStatement>()->components.size() == 0) {
+                iter = block->components.erase(iter);
+            } else {
+                iter++;
+            }
+        } else {
+            iter++;
+        }
+    }
+    return false;
+}
+
+const IR::AssignmentStatement *
+Predication::ExpressionReplacer::preorder(IR::AssignmentStatement * statement) {
+    // fill the conditions from context
+    while (context != nullptr) {
+        if (context->node->is<IR::IfStatement>()) {
+            conditions.push_back(context->node->to<IR::IfStatement>()->condition);
+        }
+        context = context->parent;
+    }
+
+    if (!statement->right->is<IR::Mux>()) {
+        statement->right = new IR::Mux(conditions.back(), statement->left, statement->left);
+    }
+    visit(statement->right);
+    return statement;
+}
+
+void Predication::ExpressionReplacer::emplaceExpression(IR::Mux * mux) {
+    auto condition = conditions[conditions.size() - currentNestingLevel];
+    bool thenElsePass = travesalPath[currentNestingLevel - 1];
+    mux->e0 = condition;
+    if (thenElsePass) {
+        mux->e1 = rightExpression;
+    } else {
+        mux->e2 = rightExpression;
+    }
+}
+
+void Predication::ExpressionReplacer::visitThen(IR::Mux * mux) {
+    auto condition = conditions[conditions.size() - currentNestingLevel];
+    auto statement = findContext<IR::AssignmentStatement>();
+    auto leftName = expr_name(statement->left);
+    auto thenExprName = expr_name(mux->e1);
+    if (expr_name(mux->e2) == expr_name(statement->left)) {
+        mux->e2 = statement->left;
+    }
+    if (mux->e1->is<IR::Mux>() || thenExprName.isNullOrEmpty() || thenExprName == leftName) {
+        if (!mux->e1->is<IR::Mux>()) {
+            mux->e1 = new IR::Mux(condition, statement->left, statement->left);
+        }
+        visit(mux->e1);
+    }
+}
+
+void Predication::ExpressionReplacer::visitElse(IR::Mux * mux) {
+    auto condition = conditions[conditions.size() - currentNestingLevel];
+    auto statement = findContext<IR::AssignmentStatement>();
+    auto leftName = expr_name(statement->left);
+    auto elseExprName = expr_name(mux->e2);
+    if (expr_name(mux->e1) == leftName) {
+        mux->e1 = statement->left;
+            mux->e1 = statement->left;
+        mux->e1 = statement->left;
+    }
+    if (mux->e2->is<IR::Mux>() || elseExprName.isNullOrEmpty() || elseExprName == leftName) {
+        if (!mux->e2->is<IR::Mux>()) {
+            mux->e2 = new IR::Mux(condition, statement->left, statement->left);
+        }
+        visit(mux->e2);
+    }
+}
+
+
+const IR::Mux * Predication::ExpressionReplacer::preorder(IR::Mux * mux) {
+    ++currentNestingLevel;
+    bool thenElsePass = travesalPath[currentNestingLevel - 1];
+    if (currentNestingLevel == travesalPath.size()) {
+        emplaceExpression(mux);
+    } else if (thenElsePass) {
+        visitThen(mux);
+    } else {
+        visitElse(mux);
+    }
+    --currentNestingLevel;
+    return mux;
+}
 
 const IR::Expression* Predication::clone(const IR::Expression* expression) {
     // We often need to clone expressions.  This is necessary because
@@ -28,13 +137,58 @@ const IR::Expression* Predication::clone(const IR::Expression* expression) {
     return expression->apply(cloner);
 }
 
-const IR::Node* Predication::postorder(IR::AssignmentStatement* statement) {
+const IR::Node* Predication::clone(const IR::AssignmentStatement* statement) {
+    // We often need to clone assignments.  This is necessary because
+    // in the end we will generate different code for the different clones of
+    // an assignments.
+    ClonePathExpressions cloner;
+    return statement->apply(cloner);
+}
+
+const IR::Node* Predication::preorder(IR::AssignmentStatement* statement) {
     if (!inside_action || ifNestingLevel == 0)
         return statement;
+    ExpressionReplacer replacer(clone(statement->right), travesalPath, getContext());
+    // Referenced lvalue is not appeard before
+    dependencies.clear();
+    visit(statement->right);
+    // print out dependencies
+    for (auto dependency : dependencies) {
+        if (liveAssignments.find(dependency) != liveAssignments.end()) {
+            // print out dependecy
+            currentBlock->push_back(liveAssignments[dependency]);
+            // remove from names to not duplicate
+            orderedNames.erase(dependency);
+            liveAssignments.erase(dependency);
+        }
+    }
+    auto statementName = expr_name(statement->left);
+    auto foundedAssignment = liveAssignments.find(statementName);
+    if (foundedAssignment != liveAssignments.end()) {
+        statement->right = foundedAssignment->second->right;
+        // move the lvalue assignment to the back
+        orderedNames.erase(statementName);
+    }
+    orderedNames.push_back(statementName);
+    auto updatedStatement = statement->apply(replacer);
+    auto rightStatement = clone(updatedStatement->to<IR::AssignmentStatement>()->right);
+    liveAssignments[statementName] = new IR::AssignmentStatement(statement->left, rightStatement);
+    return new IR::EmptyStatement();
+}
 
-    auto right = new IR::Mux(predicate(), statement->right, clone(statement->left));
-    statement->right = right;
-    return statement;
+const IR::Node* Predication::preorder(IR::PathExpression * pathExpr) {
+    dependencies.push_back(expr_name(pathExpr));
+    return pathExpr;
+}
+const IR::Node* Predication::preorder(IR::Member * member) {
+    visit(member->expr);
+    dependencies.push_back(expr_name(member));
+    return member;
+}
+const IR::Node* Predication::preorder(IR::ArrayIndex * arrInd) {
+    visit(arrInd->left);
+    dependencies.push_back(expr_name(arrInd));
+    return arrInd;
 }
 
 const IR::Node* Predication::preorder(IR::IfStatement* statement) {
@@ -43,59 +197,29 @@ const IR::Node* Predication::preorder(IR::IfStatement* statement) {
 
     ++ifNestingLevel;
     auto rv = new IR::BlockStatement;
-    cstring conditionName = generator->newName("cond");
-    auto condDecl = new IR::Declaration_Variable(conditionName, IR::Type::Boolean::get());
-    rv->push_back(condDecl);
-    auto condition = new IR::PathExpression(IR::ID(conditionName));
-
-    // A vector for a new BlockStatement.
-    auto block = new IR::BlockStatement;
-
-    const IR::Expression* previousPredicate = predicate();  // This may be nullptr
-    // a new name for the new predicate
-    cstring newPredName = generator->newName("pred");
-    predicateName.push_back(newPredName);
-    auto decl = new IR::Declaration_Variable(newPredName, IR::Type::Boolean::get());
-    block->push_back(decl);
-    // This evaluates the if condition.
-    // We are careful not to evaluate any conditional more times
-    // than in the original program, since the evaluation may have side-effects.
-    auto trueCond = new IR::AssignmentStatement(clone(condition), statement->condition);
-    block->push_back(trueCond);
-
-    const IR::Expression* pred;
-    if (previousPredicate == nullptr) {
-        pred = clone(condition);
-    } else {
-        pred = new IR::LAnd(previousPredicate, clone(condition));
-    }
-    auto truePred = new IR::AssignmentStatement(predicate(), pred);
-    block->push_back(truePred);
-
+    currentBlock = rv;
+    travesalPath.push_back(true);
     visit(statement->ifTrue);
-    block->push_back(statement->ifTrue);
+    rv->push_back(statement->ifTrue);
 
+    // This evaluates else branch
     if (statement->ifFalse != nullptr) {
-        auto neg = new IR::LNot(clone(condition));
-        auto falseCond = new IR::AssignmentStatement(clone(condition), neg);
-        block->push_back(falseCond);
-        if (previousPredicate == nullptr) {
-            pred = clone(condition);
-        } else {
-            pred = new IR::LAnd(clone(previousPredicate), clone(condition));
-        }
-        auto falsePred = new IR::AssignmentStatement(predicate(), pred);
-        block->push_back(falsePred);
-
+        travesalPath.back() = false;
         visit(statement->ifFalse);
-        block->push_back(statement->ifFalse);
+        rv->push_back(statement->ifFalse);
     }
 
-    rv->push_back(block);
-    predicateName.pop_back();
+    for (auto exprName : orderedNames) {
+        rv->push_back(liveAssignments[exprName]);
+    }
+    liveAssignments.clear();
+    orderedNames.clear();
+    travesalPath.pop_back();
     --ifNestingLevel;
+
     prune();
-    return rv;
+
+    return rv->apply(remover);
 }
 
 const IR::Node* Predication::preorder(IR::P4Action* action) {

midend/predication.h

@@ -23,51 +23,67 @@ limitations under the License.
 namespace P4 {
 
 /**
-
 This pass operates on action bodies.  It converts 'if' statements to
 '?:' expressions, if possible.  Otherwise this pass will signal an
 error.  This pass should be used only on architectures that do not
 support conditionals in actions.
-
 For this to work all statements must be assignments or other ifs.
-
 if (e)
    a = f(b);
 else
    c = f(d);
-x = y;
-
 becomes (actual implementatation is slightly optimized):
-
-bool cond;
-bool predicate = true;
 {
-  bool predicate1;
-  cond = e;
-  predicate1 = predicate && cond;
-  a = predicate1 ? f(b) : a;
-  cond = !cond;
-  predicate1 = predicate && cond;
-  c = predicate ? f(d) : c;
+    a = e ? f(b) : a;
+    c = e ? c : f(d);
 }
-x = predicate ? y : x;
-
-Not the most efficient conversion currently.
-This could be made better by looking on both the "then" and "else"
-branches, but in this way we cannot have two side-effects in the same
-conditional statement.
-
 */
 class Predication final : public Transform {
-    NameGenerator* generator;
+    class EmptyStatementRemover final : public Modifier {
+     public:
+        EmptyStatementRemover() {}
+        bool preorder(IR::BlockStatement * block) override;
+    };
+
+    /**
+     * Private Transformer only for Predication pass.
+     * This pass operates on nested Mux expressions(?:). 
+     * It replaces then and else expressions in Mux with 
+     * the appropriate expression from assignment.
+     */
+    class ExpressionReplacer final : public Transform {
+     private:
+        const IR::Expression * rightExpression;
+        const std::vector<bool>& travesalPath;
+        const Visitor::Context * context;
+        std::vector<const IR::Expression*> conditions;
+        unsigned currentNestingLevel = 0;
+     public:
+        explicit ExpressionReplacer(const IR::Expression * e,
+                                    std::vector<bool>& t,
+                                    const Visitor::Context * c)
+        : rightExpression(e), travesalPath(t), context(c)
+        { CHECK_NULL(e); }
+        const IR::Mux * preorder(IR::Mux * mux) override;
+        const IR::AssignmentStatement * preorder(IR::AssignmentStatement * statement) override;
+        void emplaceExpression(IR::Mux * mux);
+        void visitThen(IR::Mux * mux);
+        void visitElse(IR::Mux * mux);
+    };
+
+    EmptyStatementRemover remover;
+    IR::BlockStatement * currentBlock;
     bool inside_action;
-    std::vector<cstring> predicateName;
     unsigned ifNestingLevel;
+    // Traverse path of nested if-else statements
+    // true at the end of the vector means that you are currently visiting 'then' branch'
+    // false at the end of the vector means that you are in the else branch of the if statement.
+    // Size of this vector is the current if nesting level.
+    std::vector<bool> travesalPath;
+    ordered_set<cstring> orderedNames;
+    std::vector<cstring> dependencies;
+    std::map<cstring, const IR::AssignmentStatement *> liveAssignments;
 
-    const IR::Expression* predicate() const {
-        if (predicateName.empty())
-            return nullptr;
-        return new IR::PathExpression(IR::ID(predicateName.back())); }
     const IR::Statement* error(const IR::Statement* statement) const {
         if (inside_action && ifNestingLevel > 0)
             ::error(ErrorType::ERR_UNSUPPORTED_ON_TARGET,
@@ -77,15 +93,19 @@ class Predication final : public Transform {
     }
 
  public:
-    explicit Predication(NameGenerator* generator) : generator(generator),
-            inside_action(false), ifNestingLevel(0)
+    explicit Predication(NameGenerator* generator) :
+        inside_action(false), ifNestingLevel(0)
     { CHECK_NULL(generator); setName("Predication"); }
-
     const IR::Expression* clone(const IR::Expression* expression);
+    const IR::Node* clone(const IR::AssignmentStatement* statement);
     const IR::Node* preorder(IR::IfStatement* statement) override;
     const IR::Node* preorder(IR::P4Action* action) override;
     const IR::Node* postorder(IR::P4Action* action) override;
-    const IR::Node* postorder(IR::AssignmentStatement* statement) override;
+    const IR::Node* preorder(IR::AssignmentStatement* statement) override;
+    // Assignment dependecy checkers
+    const IR::Node* preorder(IR::PathExpression* pathExpr) override;
+    const IR::Node* preorder(IR::Member* member) override;
+    const IR::Node* preorder(IR::ArrayIndex* arrInd) override;
     // The presence of other statements makes predication impossible to apply
     const IR::Node* postorder(IR::MethodCallStatement* statement) override
     { return error(statement); }