diff --git a/test/Translate/Ops/pred.p4 b/test/Translate/Ops/pred.p4
new file mode 100644
index 0000000..f7ce3a7
--- /dev/null
+++ 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
})
+// CHECK: %[[TMP_1:.*]] = p4hir.alloca !p4hir.bool ["tmp_1"] : !p4hir.ref
+// CHECK: %[[TMP_2:.*]] = p4hir.alloca !p4hir.bool ["tmp_2"] : !p4hir.ref
+// 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
+// CHECK: p4hir.yield %[[TMP_2_VAL]] : !p4hir.bool
+// CHECK: }, false {
+// CHECK: %[[TMP_1_VAL:.*]] = p4hir.load %[[TMP_1]] : !p4hir.ref, !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
+
+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
index 60b9fdc..1385345 100644
--- a/tools/p4mlir-translate/translate.cpp
+++ b/tools/p4mlir-translate/translate.cpp
@@ -353,6 +353,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;
@@ -703,6 +704,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(
+ getLoc(builder, mux), getValue(mux->e0),
+ [&](mlir::OpBuilder &b, mlir::Location) {
+ visit(mux->e1);
+ b.create(getEndLoc(builder, mux->e1), getValue(mux->e1));
+ },
+ [&](mlir::OpBuilder &b, mlir::Location) {
+ visit(mux->e2);
+ b.create(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);