Add missed translation of ternary operator (aka IR::Mux)

Signed-off-by: Anton Korobeynikov <[email protected]>
p4lang · Feb 10, 2025 · c8852b1 · c8852b1
1 parent 5bdc44a
commit c8852b1
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/test/Translate/Ops/pred.p4 b/test/Translate/Ops/pred.p4
@@ -0,0 +1,25 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+// Adopted from testdata/p4_16_samples/pred.p4
+// CHECK-LABEL: p4hir.func action @cond_0(%arg0: !p4hir.bool {p4hir.dir = #p4hir<dir undir>})
+// CHECK:    %[[TMP_1:.*]] = p4hir.alloca !p4hir.bool ["tmp_1"] : !p4hir.ref<!p4hir.bool>
+// CHECK:    %[[TMP_2:.*]] = p4hir.alloca !p4hir.bool ["tmp_2"] : !p4hir.ref<!p4hir.bool>
+// CHECK:    %[[NB:.*]] = p4hir.unary(not, %arg0) : !p4hir.bool
+// CHECK:    %[[MUX:.*]] = p4hir.ternary(%[[NB]], true {
+// CHECK:      %[[TMP_2_VAL:.*]] = p4hir.load %[[TMP_2]] : !p4hir.ref<!p4hir.bool>, !p4hir.bool
+// CHECK:      p4hir.yield %[[TMP_2_VAL]] : !p4hir.bool
+// CHECK:    }, false {
+// CHECK:      %[[TMP_1_VAL:.*]] = p4hir.load %[[TMP_1]] : !p4hir.ref<!p4hir.bool>, !p4hir.bool
+// CHECK:      p4hir.yield %[[TMP_1_VAL]] : !p4hir.bool
+// CHECK:    }) : (!p4hir.bool) -> !p4hir.bool
+// CHECK:    p4hir.store %[[MUX]], %[[TMP_1]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+
+action cond_0(bool in_b) {
+   bit<32> a;
+   bool tmp_1;
+   bool tmp_2;
+
+   tmp_1 = (!in_b ? tmp_2 : tmp_1);
+   tmp_2 = (!in_b && !!in_b ? a == 32w5 : tmp_2);
+   tmp_1 = (!in_b && !!in_b ? (!in_b && !!in_b ? a == 32w5 : tmp_2) : (!in_b && !in_b ? false : tmp_1));
+}
diff --git a/tools/p4mlir-translate/translate.cpp b/tools/p4mlir-translate/translate.cpp
@@ -337,6 +337,7 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
 
     bool preorder(const P4::IR::Declaration_Constant *decl) override;
     bool preorder(const P4::IR::AssignmentStatement *assign) override;
+    bool preorder(const P4::IR::Mux *mux) override;
     bool preorder(const P4::IR::LOr *lor) override;
     bool preorder(const P4::IR::LAnd *land) override;
     bool preorder(const P4::IR::IfStatement *ifs) override;
@@ -651,6 +652,29 @@ bool P4HIRConverter::preorder(const P4::IR::LAnd *land) {
     return false;
 }
 
+bool P4HIRConverter::preorder(const P4::IR::Mux *mux) {
+    ConversionTracer trace("Converting ", mux);
+
+    // Materialize condition first
+    visit(mux->e0);
+
+    // Make the value itself
+    auto value = builder.create<P4HIR::TernaryOp>(
+        getLoc(builder, mux), getValue(mux->e0),
+        [&](mlir::OpBuilder &b, mlir::Location) {
+            visit(mux->e1);
+            b.create<P4HIR::YieldOp>(getEndLoc(builder, mux->e1), getValue(mux->e1));
+        },
+        [&](mlir::OpBuilder &b, mlir::Location) {
+            visit(mux->e2);
+            b.create<P4HIR::YieldOp>(getEndLoc(builder, mux->e2), getValue(mux->e2));
+        });
+
+    setValue(mux, value.getResult());
+
+    return false;
+}
+
 bool P4HIRConverter::preorder(const P4::IR::IfStatement *ifs) {
     ConversionTracer trace("Converting ", ifs);