Add switch elimination pass for BMv2 target support Transforms switch… by devalgupta404 · Pull Request #298 · p4lang/p4mlir-incubator

devalgupta404 · 2026-01-05T20:01:25Z

Add switch elimination pass for BMv2 target support

Closes #252
Parent issue: #247

Summary

Transforms switch statements on arbitrary expressions to table-based switches on action_run. Required for hardware targets (like BMv2) that don't support arbitrary expression switches.

Creates a temporary variable to store the switch condition expression before table apply, following p4c's pattern.

Reference

Based on p4c's midend implementation:

https://github.com/p4lang/p4c/blob/main/midend/eliminateSwitch.cpp

Transformation

Input (switch on bit<32>):

p4hir.switch (%val : !p4hir.bit<32>) {
  p4hir.case(anyof, [#int16, #int32]) { assign 1 }
  p4hir.case(equal, [#int2]) { assign 2 }
  p4hir.case(default, []) { assign 3 }
}

Output (table + switch on action_run):

p4hir.func action @_switch_0_case_0() annotations {hidden} { return }
p4hir.func action @_switch_0_case_1() annotations {hidden} { return }
p4hir.func action @_switch_0_default() annotations {hidden} { return }

p4hir.table @_switch_0_table annotations {hidden} {
  p4hir.table_key: exact match on condition
  p4hir.table_entries: 16,32 → case_0, 2 → case_1
  p4hir.table_default_action: _switch_0_default
}

// Temporary variable for switch key (following p4c's pattern)
%_switch_0_key = p4hir.variable ["_switch_0_key"] annotations {hidden}
p4hir.assign %val, %_switch_0_key
%key_val = p4hir.read %_switch_0_key

%result = p4hir.table_apply @_switch_0_table with key(%key_val)
%action_run = p4hir.struct_extract %result["action_run"]
p4hir.switch (%action_run : !action_enum) { ... }

Test Evidence

$ p4mlir-opt --p4hir-switch-elimination test/Transforms/Passes/switch-elimination.mlir | FileCheck test/Transforms/Passes/switch-elimination.mlir
FileCheck PASSED

Full pass output (click to expand)

!_switch_0_action_enum = !p4hir.enum<"_switch_0_action_enum", _switch_0_case_0, _switch_0_case_1, _switch_0_default>
!b32i = !p4hir.bit<32>
!_switch_0_result = !p4hir.struct<"_switch_0_result", hit: !p4hir.bool, miss: !p4hir.bool, action_run: !_switch_0_action_enum>

module {
  p4hir.control @c(%arg0: !p4hir.ref<!b32i>)() {
    p4hir.func action @_switch_0_case_0() annotations {hidden} {
      p4hir.return
    }
    p4hir.func action @_switch_0_case_1() annotations {hidden} {
      p4hir.return
    }
    p4hir.func action @_switch_0_default() annotations {hidden} {
      p4hir.return
    }
    p4hir.table @_switch_0_table annotations {hidden} {
      p4hir.table_key(%arg1: !b32i) {
        p4hir.match_key #exact %arg2 : !b32i
      }
      p4hir.table_actions {
        p4hir.table_action @_switch_0_case_0() { ... }
        p4hir.table_action @_switch_0_case_1() { ... }
        p4hir.table_action @_switch_0_default() { ... }
      }
      p4hir.table_default_action {
        p4hir.call @c::@_switch_0_default () : () -> ()
      }
      p4hir.table_entries const {
        p4hir.table_entry [16] -> case_0
        p4hir.table_entry [32] -> case_0
        p4hir.table_entry [2] -> case_1
      }
    }
    p4hir.control_apply {
      %val = p4hir.read %arg0 : <!b32i>
      %_switch_0_key = p4hir.variable ["_switch_0_key"] annotations {hidden} : <!b32i>
      p4hir.assign %val, %_switch_0_key : <!b32i>
      %val_0 = p4hir.read %_switch_0_key : <!b32i>
      %result = p4hir.table_apply @c::@_switch_0_table with key(%val_0)
      %action_run = p4hir.struct_extract %result["action_run"]
      p4hir.switch (%action_run : !_switch_0_action_enum) {
        p4hir.case(equal, [case_0]) { assign 1 }
        p4hir.case(equal, [case_1]) { assign 2 }
        p4hir.case(default, [default]) { assign 3 }
      }
    }
  }
}

Files Changed

File	Description
`lib/Transforms/SwitchElimination.cpp`	Pass implementation
`include/p4mlir/Transforms/Passes.td`	Pass definition
`include/p4mlir/Transforms/Passes.h`	Pass declaration
`lib/Transforms/CMakeLists.txt`	Build integration
`test/Transforms/Passes/switch-elimination.mlir`	LIT test

How to Test

# Build
ninja extensions/p4mlir/bin/p4mlir-opt

# Run test
p4mlir-opt --p4hir-switch-elimination test/Transforms/Passes/switch-elimination.mlir

# Run all tests
ninja check-p4mlir

devalgupta404 · 2026-01-06T16:27:44Z

@PietroGhg please review it

PietroGhg · 2026-01-07T11:44:34Z

Hi @devalgupta404 thanks for this PR, I haven't looked in detail at the code yet, just from looking at the output of the test you added I think that this is missing the creation of a temporary variable to which we assign the condition expression of the switch statement (it's done here in the original p4c pass). I think this is necessary to do otherwise the whole apply block can't be correctly expressed in the final json representation. Please add that and then I'll take a better look at the code, thanks :)

PietroGhg · 2026-01-08T10:46:43Z

lib/Transforms/SwitchElimination.cpp

+
+namespace {
+
+class NameGenerator {


Could you use SymbolTable::generateSymbolName for this, e.g. translate.cpp

PietroGhg · 2026-01-08T10:47:16Z

lib/Transforms/SwitchElimination.cpp

+    void runOnOperation() override;
+
+private:
+    LogicalResult eliminateSwitch(P4HIR::ControlOp control, P4HIR::SwitchOp switchOp,


nit: eliminateSwitch and isActionRunSwitch can be made static

PietroGhg · 2026-01-08T11:36:38Z

lib/Transforms/SwitchElimination.cpp

+            {
+                OpBuilder::InsertionGuard guard(tableBuilder);
+                tableBuilder.setInsertionPointToStart(&tableKeyOp.getBody().front());
+                auto keyArg = tableKeyOp.getBody().front().addArgument(condType, tableLoc);


I think the argument here is already present and you should just use getArgument(0)

mtsamis

Hi devalgupta404, thanks for the PR!

I've made some comments for an issue with isActionRunSwitch plus some others for the code structure, please have a look.

mtsamis · 2026-01-08T12:17:28Z

lib/Transforms/SwitchElimination.cpp

+    bool isActionRunSwitch(P4HIR::SwitchOp switchOp);
+};
+
+bool SwitchEliminationPass::isActionRunSwitch(P4HIR::SwitchOp switchOp) {


This doesn't look enough to classify a switch as an action_run one. You're looking if the type of the condition is a struct which has an action_run field which has two issues:

You don't know if action_run is actually used in that particular switch, only that it exists in general.

You cannot differentiate from a user-defined type that has an unrelated field named action_run (since it isn't really a unique name).

I think you need to inspect the condition value rather than the type. This way you should be able to address these issues by e.g. following it's definitions and make sure it's coming from a table_apply, plus that action_run is the actual field used there.

mtsamis · 2026-01-08T12:24:13Z

lib/Transforms/SwitchElimination.cpp

+    }
+
+    llvm::SmallVector<std::string> actionNames;
+    llvm::SmallVector<mlir::Attribute> enumFields;


It looks from below that enumFields is always builder.getStringAttr(actionNames), so not necessary to keep two vectors, just one.

Could probably be a single vector of mlir::StringAttr (so you can get the string as needed). If it's only used in FlatSymbolRefAttr you can probably just pass the StringAttr there directly anyway.

mtsamis · 2026-01-08T12:26:41Z

lib/Transforms/SwitchElimination.cpp

+            tableBuilder.create<P4HIR::TableEntriesOp>(
+                tableLoc, true, DictionaryAttr(),
+                [&](OpBuilder &entriesBuilder, Location entriesLoc) {
+                    for (size_t i = 0; i < cases.size(); ++i) {


Can you please replace all loops that have manual iteration with range-based ones? If you need the index you can use llvm::enumerate, but e.g. here it's not really needed, it's better to do a loop with llvm::zip(cases, actionsNames)

mtsamis · 2026-01-08T12:29:48Z

lib/Transforms/SwitchElimination.cpp

+                        for (auto valueAttr : caseOp.getValue()) {
+                            auto tupleType = TupleType::get(ctx, {condType});
+                            auto keyAttr = P4HIR::AggAttr::get(
+                                tupleType, builder.getArrayAttr({valueAttr}));


Better avoid mixing builder and entriesBuilder here (and in some other places).

Also maybe you'd want to consider naming the lambda builder arguments something sort like just b (consistent with other places in the codebase, so I think it's fine). It may help to reduce line length for readability.

mtsamis · 2026-01-08T12:31:08Z

lib/Transforms/SwitchElimination.cpp

+    // Create a temporary variable for the switch key (following p4c's pattern)
+    std::string keyName = prefix + "_key";
+    auto keyVar = builder.create<P4HIR::VariableOp>(
+        loc, P4HIR::ReferenceType::get(ctx, condType), keyName, hiddenAttr);


ctx in P4HIR::ReferenceType::get(ctx, condType) is redundant, you can use P4HIR::ReferenceType::get(condType) (same for some other places).

mtsamis · 2026-01-08T12:33:17Z

lib/Transforms/SwitchElimination.cpp

+                            if (!mlir::isa<P4HIR::YieldOp>(&op))
+                                caseBuilder.clone(op, mapper);
+                        }
+                        caseBuilder.create<P4HIR::YieldOp>(caseLoc);


Same as above

mtsamis · 2026-01-08T12:48:54Z

lib/Transforms/SwitchElimination.cpp

+    auto hiddenAttr = builder.getDictionaryAttr(
+        {builder.getNamedAttr("hidden", builder.getUnitAttr())});
+
+    for (size_t i = 0; i < cases.size(); ++i) {


There is some fair amount of duplication with the handling of normal / default cases. E.g. here the body which is not small is duplicated below with only the name changing. You can consider using a lambda with the code outside (e.g. createActionForCase(llvm::StringRef name), to avoid the duplication.

I'm not sure, but it may also be possible to eliminate the cases vector, and iterate over cases with an if for the default case and avoid the duplication altogether.

mtsamis · 2026-01-08T12:49:19Z

lib/Transforms/SwitchElimination.cpp

+                    });
+            }
+
+            if (defaultCase) {


Same as at the top for default case duplication.

mtsamis · 2026-01-08T12:51:04Z

lib/Transforms/SwitchElimination.cpp

+    auto module = getOperation();
+    NameGenerator nameGen;
+
+    llvm::SmallVector<std::pair<P4HIR::ControlOp, P4HIR::SwitchOp>> toProcess;


It looks like we don't need a vector here, we can just call eliminateSwitch while iterating in a single loop and e.g. use WalkResult::interrupt if it fails. Also isActionRunSwitch is also checked at the top of eliminateSwitch so there's no need to double call it, you can pick one place that you like.

mtsamis · 2026-01-08T12:54:47Z

test/Transforms/Passes/switch-elimination.mlir

+
+// CHECK-LABEL: p4hir.control @c
+module {
+  p4hir.control @c(%arg0: !p4hir.ref<!b32i> {p4hir.dir = #p4hir<dir inout>, p4hir.param_name = "b"})() {


The CHECK directives here look insufficient. Can you please match the IR that we expect to be generated from this pass? Looking at your pass I believe we would at least need to match the newly introduced p4hir.table's contents and the resulting SwitchOp / actions?

Second that. We also need to ensure that whole logic is covered including handling of action_run switches.

You'd better take the existing p4c testcases and port them if you're out of ideas :)

devalgupta404 · 2026-01-11T16:21:34Z

@mtsamis @PietroGhg I’ve made the changes as suggested above. Could you please review it?

PietroGhg · 2026-01-13T09:12:24Z

Hi @devalgupta404 changes look good on my side, let's wait for @mtsamis to take a look too

devalgupta404 · 2026-01-14T09:51:58Z

@mtsamis please take a look on updated changes

mtsamis · 2026-01-14T11:13:12Z

Hi @devalgupta404, looking at the latest version of your code I can say it still needs more polishing;

In general, things I have mentioned in my previous detailed review aren't fully addressed (code duplication, elimination of cases vector, good tests with full IR matching, use of builders and their names) or were applied hastily (e.g. comment about yield copying was partially addressed, llvm::enumerate is overused when there are better alternatives). I kindly ask you to have a look at these again and try to improve the implementation based on the ideas rather than just applying my suggestions 🙂

Since I did re-review the implementation, I can also provide some further feedback to help you refining this.

On the high level:

Some parts of code are currently not readable due to being terse, nested and not appropriately commented (Your table and switch op building). These currently contain three nested lambdas, with the inner ones being quite hard to read. You don't necessarily have to nest everything together in these lambdas, you can probably just build the outer op and then add things after the fact inside it.

I will also add some inline comments, for more specific things.

mtsamis · 2026-01-14T11:14:47Z

lib/Transforms/SwitchElimination.cpp

+        [&](llvm::StringRef candidate) {
+            return symbolTable.lookup(candidate) != nullptr;
+        },
+        counter).str().str();


.str().str(); looks like code smell, probably can be refactored to something better (you can see examples of generateSymbolName in mlir repo maybe).

mtsamis · 2026-01-14T11:15:55Z

lib/Transforms/SwitchElimination.cpp

-        actionNames.push_back(actionName);
-        enumFields.push_back(builder.getStringAttr(actionName));
-
+    auto createAction = [&](mlir::StringAttr actionName) {


There's no reason to keep actionNames.push_back or creation of stringattr outside this lambda, that's why I mentioned a llvm::StringRef version in my review.

mtsamis · 2026-01-14T11:16:28Z

lib/Transforms/SwitchElimination.cpp

+        },
+        counter).str().str();

    llvm::SmallVector<P4HIR::CaseOp> cases;


I still believe that we shouldn't have this vector and refactor the rest of the code appropriately.

mtsamis · 2026-01-14T11:17:46Z

lib/Transforms/SwitchElimination.cpp

        auto funcOp = P4HIR::FuncOp::buildAction(builder, loc, actionName, funcType,
-                                                  {}, hiddenAttr);
+                                                 {}, hiddenAttr);
        OpBuilder funcBuilder(funcOp.getBody());


Let's avoid creating OpBuilders, maybe just an InsertionGuard would do here?

mtsamis · 2026-01-14T11:19:26Z

lib/Transforms/SwitchElimination.cpp


    auto loc = switchOp.getLoc();
    auto condType = switchOp.getCondition().getType();
    OpBuilder builder(switchOp);


You're currently using an OpBuilder to create ops and direct calls for things like erase. It would be better to create an IRRewriter instead of OpBuilder and do everything from there.

mtsamis · 2026-01-14T11:19:52Z

lib/Transforms/SwitchElimination.cpp

                    defaultBuilder.create<P4HIR::CallOp>(defaultLoc, controlRef, ValueRange{});
                });
-
            tableBuilder.create<P4HIR::TableEntriesOp>(


See high level comment for all this code

mtsamis · 2026-01-14T11:22:57Z

lib/Transforms/SwitchElimination.cpp

-        control->walk([&](P4HIR::SwitchOp switchOp) {
-            if (!isActionRunSwitch(switchOp))
-                toProcess.emplace_back(control, switchOp);
+        return control->walk([&](P4HIR::SwitchOp switchOp) {


I think we would need to propagate the error and do something like signalPassFailure if there's one using wasInterrupted? Also maybe we should consider using an OpRewritePattern if applicable (probably on control op)?

mtsamis · 2026-01-14T11:24:04Z

test/Transforms/Passes/switch-elimination.mlir

-      // CHECK: p4hir.table_apply @c::@_switch_0_table with key(%[[KEY_VAL]])
-      // CHECK: p4hir.struct_extract
-      // CHECK: p4hir.switch
+      // CHECK: %{{.*}} = p4hir.variable ["_switch{{.*}}_key"] annotations {hidden}


These tests are still very insufficient for what the pass is doing. Please have a look at other tests from optimization passes that do IR transformations (one example could be test/Transforms/Passes/InlineControls/inline-3.mlir)

mtsamis · 2026-01-14T11:27:31Z

lib/Transforms/SwitchElimination.cpp

+        extractOp.getInput().getDefiningOp<P4HIR::TableApplyOp>();
+    if (!tableApply)
+        return false;
+    return true;


Let's avoid unnecessarily verbose code here.

Maybe just if (!extractOp || extractOp.getFieldName() != "action_run") ...

And return extractOp.getInput().getDefiningOp<P4HIR::TableApplyOp>() != nullptr or similar

mtsamis · 2026-01-14T12:57:56Z

Also, it looks like a variable is introduced that is written and then subsequently read. This doesn't do something in p4mlir and should probably be deleted; it will be canonicalized away anyway.

PietroGhg · 2026-01-14T13:19:40Z

Also, it looks like a variable is introduced that is written and then subsequently read. This doesn't do something in p4mlir and should probably be deleted; it will be canonicalized away anyway.

I asked for that change (the variable is needed otherwise we can't represent this "pattern" in bmv2), but I didn't think about the canonicalizer removing it. I'm ok with removing it in this pass, we'll handle things in some BMv2-specific pass (sorry for the back and forth @devalgupta404)

devalgupta404 · 2026-01-20T05:49:30Z

@mtsamis i applied the requested changes please take a look.

mtsamis · 2026-01-27T23:19:36Z

@mtsamis i applied the requested changes please take a look.

Thanks! Please give me a bit more time to review this 🙂

mtsamis · 2026-02-09T09:19:09Z

Hi @devalgupta404, I went over the current implementation again. The implementation looks better now, but I see that a number of the requested changes are not implemented. Can you revisit and either implement the changes or discuss why you believe they should not be implemented? Thanks!

Transforms switch statements on arbitrary expressions to table-based switches on action_run. Creates a temporary variable to store the switch condition expression before table apply, following p4c's pattern. Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Structured bindings cannot be captured in lambdas in C++17 (only C++20+). Changed llvm::zip loops to llvm::enumerate and extract actionName separately to avoid capturing structured bindings inside nested lambdas. Fixes CI build error on Ubuntu 22.04 with clang 14. Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

…th OpBuilder::InsertionGuard on the local builder Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

devalgupta404 · 2026-02-10T14:37:01Z

Hi @mtsamis regarding llvm::zip, switchOp.cases(), pre collecting non default cases in the SmallVector just to enable zip adds unnecessary complexity so instead i use the single loop with conditional logic for default case and there are few changes i missed earlier so now i made all lambda builders renamed to b and replaced rewriter mixing with OpBuilder::InsertionGuard and also extract cloneCase helper with single loop over all cases.

mtsamis

Hi @mtsamis regarding llvm::zip, switchOp.cases(), pre collecting non default cases in the SmallVector just to enable zip adds unnecessary complexity so instead i use the single loop with conditional logic for default case and there are few changes i missed earlier so now i made all lambda builders renamed to b and replaced rewriter mixing with OpBuilder::InsertionGuard and also extract cloneCase helper with single loop over all cases.

Right, no problem with not using zip if not a good fit.

I was mostly referring to few other things. The most important things that this still needs:

More readable code: As I mentioned already, the triply-nested builder lambdas make this quite hard to read. These don't really have to be nested as far as I see, so can you refactor that please?
The tests need some improvement as well, {{.*}} is used in places where named matches should be used instead.

mtsamis · 2026-03-02T13:22:48Z

lib/Transforms/SwitchElimination.cpp

+                    [&](OpBuilder &b, Location caseLoc) {
+                        IRMapping mapper;
+                        for (auto &op : origCase.getCaseRegion().front())
+                            b.clone(op, mapper);


Can this something like inlineBlockBefore/inlineRegionBefore be used here instead?

devalgupta404 force-pushed the issue-252-switch-elimination branch from 6c83110 to 399da9b Compare January 7, 2026 21:17

PietroGhg reviewed Jan 8, 2026

View reviewed changes

mtsamis reviewed Jan 8, 2026

View reviewed changes

devalgupta404 force-pushed the issue-252-switch-elimination branch 2 times, most recently from b197d48 to 197d619 Compare January 10, 2026 06:05

devalgupta404 force-pushed the issue-252-switch-elimination branch from 197d619 to 47ea203 Compare January 12, 2026 14:06

mtsamis requested changes Jan 14, 2026

View reviewed changes

devalgupta404 force-pushed the issue-252-switch-elimination branch 10 times, most recently from b4e69be to 1ae6f5f Compare January 18, 2026 11:27

devalgupta404 added 5 commits February 10, 2026 19:56

updte as per the maintainers requirement

57b8b93

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Refactor switch elimination per maintainer feedback Address

5a5f902

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Fix build issues

4a315d6

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

Fix switch elimination test patterns - verified locally

6192d54

Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

devalgupta404 force-pushed the issue-252-switch-elimination branch from 1ae6f5f to 8de2ceb Compare February 10, 2026 14:26

Rename all lambda builder arguments to and replace rewriter mixing wi…

d826d66

…th OpBuilder::InsertionGuard on the local builder Signed-off-by: devalgupta404 <devalgupta4@gmail.com>

devalgupta404 force-pushed the issue-252-switch-elimination branch from 8de2ceb to d826d66 Compare February 10, 2026 14:29

mtsamis requested changes Mar 2, 2026

View reviewed changes


		namespace {

		class NameGenerator {

Conversation

devalgupta404 commented Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!