Pass to get rid of action_run() calls

- introduce new metadata set in the actions of the table

Signed-off-by: Chris Dodd <cdodd@nvidia.com>

backends/p4test/midend.cpp

@@ -33,6 +33,7 @@ limitations under the License.
 #include "midend/complexComparison.h"
 #include "midend/copyStructures.h"
 #include "midend/def_use.h"
+#include "midend/eliminateActionRun.h"
 #include "midend/eliminateInvalidHeaders.h"
 #include "midend/eliminateNewtype.h"
 #include "midend/eliminateSerEnums.h"
@@ -125,7 +126,7 @@ MidEnd::MidEnd(CompilerOptions &options, std::ostream *outStream) {
          new P4::SimplifyControlFlow(&typeMap),
          new P4::CompileTimeOperations(),
          new P4::TableHit(&typeMap),
-         new P4::EliminateSwitch(&typeMap),
+         options.preferSwitch ? new P4::ElimActionRun() : new P4::EliminateSwitch(&typeMap),
          new P4::ResolveReferences(&refMap),
          new P4::TypeChecking(&refMap, &typeMap, true),  // update types before ComputeDefUse
          new PassRepeated({

backends/p4test/p4test.cpp

@@ -42,6 +42,7 @@ class P4TestOptions : public CompilerOptions {
     bool parseOnly = false;
     bool validateOnly = false;
     bool loadIRFromJson = false;
+    bool preferSwitch = false;
     P4TestOptions() {
         registerOption(
             "--listMidendPasses", nullptr,
@@ -84,6 +85,13 @@ class P4TestOptions : public CompilerOptions {
             "Turn off LOGN() statements in the compiler.\n"
             "Use '@__debug' annotation to enable LOGN on "
             "the annotated P4 object within the source code.\n");
+        registerOption(
+            "--preferSwitch", nullptr,
+            [this](const char *) {
+                preferSwitch = true;
+                return true;
+            },
+            "use passes that convert to switch instead eliminating them");
     }
 };
 

midend/CMakeLists.txt

@@ -21,6 +21,7 @@ set (MIDEND_SRCS
   copyStructures.cpp
   coverage.cpp
   def_use.cpp
+  eliminateActionRun.cpp
   eliminateInvalidHeaders.cpp
   eliminateNewtype.cpp
   eliminateSerEnums.cpp
@@ -74,6 +75,7 @@ set (MIDEND_HDRS
   copyStructures.h
   coverage.h
   def_use.h
+  eliminateActionRun.h
   eliminateInvalidHeaders.h
   eliminateNewtype.h
   eliminateSerEnums.h

midend/eliminateActionRun.cpp

@@ -0,0 +1,124 @@
+/*
+Copyright 2024 Intel Corp.
+
+Licensed 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.
+*/
+
+#include "eliminateActionRun.h"
+
+#include "frontends/p4/methodInstance.h"
+
+namespace P4 {
+
+void ElimActionRun::ActionTableUse::postorder(const IR::Member *mem) {
+    if (mem->member != "action_run") return;
+    auto *mce = mem->expr->to<IR::MethodCallExpression>();
+    if (!mce) return;
+    auto *mi = MethodInstance::resolve(mce, this);
+    auto *tbl = mi->object->to<IR::P4Table>();
+    if (!tbl) return;
+    LOG3("found: " << mem);
+    auto [it, inserted] = self.actionRunTables.emplace(tbl->name, tbl);
+    if (!inserted) {
+        LOG3("  (additional apply() of table)");
+        return;
+    }
+    auto *info = &it->second;
+    IR::IndexedVector<IR::Declaration_ID> tags;
+    for (auto *ale : tbl->getActionList()->actionList) {
+        cstring name = ale->getName();
+        BUG_CHECK(!self.actionsToModify.count(name),
+                  "action %s used in multiple tables (%s and %s) -- LocalizeActions must be run "
+                  "before ElimActionRun",
+                  name, self.actionsToModify.at(name)->tbl->name, tbl->name);
+        self.actionsToModify[name] = info;
+        cstring tag = self.nameGen.newName(tbl->name + "_" + name);
+        info->actions.emplace(name, tag);
+        tags.push_back(new IR::Declaration_ID(tag));
+        LOG4("  action " << name << " [" << ale->expression << "] -> " << tag);
+    }
+    info->action_tags =
+        new IR::Type_Enum(self.nameGen.newName(tbl->name + "_action_run_t"), std::move(tags));
+    info->action_run = new IR::Declaration_Variable(self.nameGen.newName(tbl->name + "_action_run"),
+                                                    new IR::Type_Name(info->action_tags->name));
+}
+
+const IR::Expression *ElimActionRun::RewriteActionRun::postorder(IR::Member *mem) {
+    if (mem->member != "action_run") return mem;
+    auto *mce = mem->expr->to<IR::MethodCallExpression>();
+    if (!mce) return mem;
+    auto *mi = MethodInstance::resolve(mce, this);
+    auto *tbl = mi->object->to<IR::P4Table>();
+    if (!tbl) return mem;
+    auto &info = self.actionRunTables.at(tbl->name);
+    auto *parent = findContext<IR::Statement>();
+    BUG_CHECK(parent && parent->is<IR::SwitchStatement>(), "action_run not in switch");
+    auto &pps = self.prepend_statement[parent];
+    if (!pps) pps = new IR::Vector<IR::StatOrDecl>;
+    pps->push_back(new IR::MethodCallStatement(mce));
+    self.add_to_control.push_back(info.action_run);
+    return new IR::PathExpression(new IR::Type_Name(info.action_tags->name), info.action_run->name);
+}
+
+const IR::P4Action *ElimActionRun::RewriteActionRun::preorder(IR::P4Action *act) {
+    auto *info = get(self.actionsToModify, act->name);
+    if (!info) return act;
+    auto *body = act->body->clone();
+    auto *tag_type = new IR::Type_Name(info->action_tags->name);
+
+    body->components.insert(
+        body->components.begin(),
+        new IR::AssignmentStatement(
+            new IR::PathExpression(tag_type->clone(), info->action_run->name),
+            new IR::Member(new IR::TypeNameExpression(tag_type), info->actions.at(act->name))));
+    act->body = body;
+    return act;
+}
+
+const IR::SwitchCase *ElimActionRun::RewriteActionRun::postorder(IR::SwitchCase *swCase) {
+    auto *swtch = getParent<IR::SwitchStatement>();
+    BUG_CHECK(swtch, "case not in switch");
+    if (!self.prepend_statement.count(swtch)) return swCase;
+    if (swCase->label->is<IR::DefaultExpression>()) return swCase;
+    auto *pe = swCase->label->to<IR::PathExpression>();
+    BUG_CHECK(pe, "case label is not an action name");
+    auto *info = self.actionsToModify.at(pe->path->name);
+    auto *tag_type = new IR::Type_Name(info->action_tags->name);
+    swCase->label =
+        new IR::Member(new IR::TypeNameExpression(tag_type), info->actions.at(pe->path->name));
+    return swCase;
+}
+
+const IR::Node *ElimActionRun::RewriteActionRun::postorder(IR::Statement *stmt) {
+    auto *pps = get(self.prepend_statement, stmt);
+    if (!pps) return stmt;
+    pps->push_back(stmt);
+    self.prepend_statement.erase(stmt);
+    return pps;
+}
+
+const IR::P4Control *ElimActionRun::RewriteActionRun::postorder(IR::P4Control *ctrl) {
+    BUG_CHECK(self.prepend_statement.empty(), "orphan action_run in control %s", ctrl->name);
+    ctrl->controlLocals.prepend(self.add_to_control);
+    if (!self.add_to_control.empty()) LOG4(ctrl);
+    self.add_to_control.clear();
+    return ctrl;
+}
+
+const IR::P4Program *ElimActionRun::RewriteActionRun::postorder(IR::P4Program *prog) {
+    auto front = prog->objects.begin();
+    for (auto &[_, info] : self.actionRunTables) prog->objects.insert(front, info.action_tags);
+    return prog;
+}
+
+}  // namespace P4

midend/eliminateActionRun.h

@@ -0,0 +1,87 @@
+/*
+Copyright 2024 Intel Corp.
+
+Licensed 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.
+*/
+
+#ifndef MIDEND_ELIMINATEACTIONRUN_H_
+#define MIDEND_ELIMINATEACTIONRUN_H_
+
+#include "frontends/common/resolveReferences/resolveReferences.h"
+#include "ir/ir.h"
+
+namespace P4 {
+
+class ElimActionRun : public PassManager {
+    MinimalNameGenerator nameGen;
+
+    struct atinfo_t {
+        const IR::P4Table *tbl;
+        const IR::Type_Enum *action_tags;
+        const IR::Declaration_Variable *action_run;
+        std::map<cstring, cstring> actions;
+        explicit atinfo_t(const IR::P4Table *t) : tbl(t) {}
+    };
+
+    std::map<cstring, atinfo_t> actionRunTables;
+    std::map<cstring, atinfo_t *> actionsToModify;
+    std::map<const IR::Statement *, IR::Vector<IR::StatOrDecl> *> prepend_statement;
+    IR::Vector<IR::Declaration> add_to_control;
+
+    class ActionTableUse : public Inspector, public ResolutionContext {
+        ElimActionRun &self;
+
+     public:
+        explicit ActionTableUse(ElimActionRun &s) : self(s) {}
+
+        bool preorder(const IR::P4Parser *) override { return false; }
+        bool preorder(const IR::P4Action *) override { return false; }
+        bool preorder(const IR::Function *) override { return false; }
+
+        void postorder(const IR::Member *) override;
+    } actionTableUse;
+
+    class RewriteActionRun : public Transform, public ResolutionContext {
+        ElimActionRun &self;
+
+     public:
+        explicit RewriteActionRun(ElimActionRun &s) : self(s) {}
+
+        const IR::P4Parser *preorder(IR::P4Parser *p) override {
+            prune();
+            return p;
+        }
+        const IR::Function *preorder(IR::Function *f) override {
+            prune();
+            return f;
+        }
+
+        const IR::Expression *postorder(IR::Member *) override;
+        const IR::P4Action *preorder(IR::P4Action *) override;
+        const IR::SwitchCase *postorder(IR::SwitchCase *) override;
+        const IR::Node *postorder(IR::Statement *) override;
+        const IR::P4Control *postorder(IR::P4Control *) override;
+        const IR::P4Program *postorder(IR::P4Program *) override;
+    };
+
+ public:
+    ElimActionRun() : actionTableUse(*this) {
+        addPasses({&actionTableUse, new PassIf([this]() { return !actionRunTables.empty(); },
+                                               {&nameGen, new RewriteActionRun(*this),
+                                                [this]() { actionRunTables.clear(); }})});
+    }
+};
+
+}  // namespace P4
+
+#endif /* MIDEND_ELIMINATEACTIONRUN_H_ */