Add if operator

Signed-off-by: Anton Korobeynikov <[email protected]>

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.h

@@ -15,6 +15,11 @@
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "p4mlir/Dialect/P4HIR/P4HIR_OpsEnums.h"
 #include "p4mlir/Dialect/P4HIR/P4HIR_Types.h"
+
+namespace P4::P4MLIR::P4HIR {
+void buildTerminatedBody(mlir::OpBuilder &builder, mlir::Location loc);
+}  // namespace  P4::P4MLIR::P4HIR
+
 #define GET_OP_CLASSES
 #include "p4mlir/Dialect/P4HIR/P4HIR_Ops.h.inc"
 

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -381,8 +381,7 @@ def ScopeOp : P4HIR_Op<"scope", [
 }
 
 def YieldOp : P4HIR_Op<"yield", [ReturnLike, Terminator,
-    ParentOneOf<["ScopeOp", "TernaryOp",
-                 // "IfOp",
+    ParentOneOf<["ScopeOp", "TernaryOp", "IfOp",
                  // "SwitchOp", "CaseOp",
                  // "ForInOp", "ForOp",
                  // "CallOp"
@@ -490,5 +489,47 @@ def TernaryOp : P4HIR_Op<"ternary",
   }];
 }
 
+def IfOp : P4HIR_Op<"if",
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+       RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments]> {
+  let summary = "The if-then-else operation";
+  let description = [{
+    The `p4hir.if` operation represents an if-then-else construct for
+    conditionally executing two regions of code. The operand is a `p4hir.bool`
+    type.
+
+    Examples:
+
+    ```mlir
+    p4hir.if %b  {
+      ...
+    } else {
+      ...
+    }
+
+    p4hir.if %c  {
+      ...
+    }
+    ```
+
+    `p4hir.if` defines no values and the 'else' can be omitted. The if/else
+    regions must be terminated. If the region has only one block, the terminator
+    can be left out, and `p4hir.yield` terminator will be inserted implictly.
+    Otherwise, the region must be explicitly terminated.
+  }];
+  let arguments = (ins BooleanType:$condition);
+  let regions = (region AnyRegion:$thenRegion, AnyRegion:$elseRegion);
+
+  let hasCustomAssemblyFormat = 1;
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<(ins "mlir::Value":$cond, "bool":$withElseRegion,
+      CArg<"llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)>",
+           "P4HIR::buildTerminatedBody">:$thenBuilder,
+      CArg<"llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)>",
+           "nullptr">:$elseBuilder)>
+  ];
+}
 
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_OPS_TD

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -206,6 +206,103 @@ void P4HIR::TernaryOp::build(OpBuilder &builder, OperationState &result, Value c
     if (yield.getNumOperands() == 1) result.addTypes(TypeRange{yield.getOperandTypes().front()});
 }
 
+//===----------------------------------------------------------------------===//
+// IfOp
+//===----------------------------------------------------------------------===//
+
+ParseResult P4HIR::IfOp::parse(OpAsmParser &parser, OperationState &result) {
+    // Create the regions for 'then'.
+    result.regions.reserve(2);
+    Region *thenRegion = result.addRegion();
+    Region *elseRegion = result.addRegion();
+
+    auto &builder = parser.getBuilder();
+    OpAsmParser::UnresolvedOperand cond;
+    Type boolType = P4HIR::BoolType::get(builder.getContext());
+
+    if (parser.parseOperand(cond) || parser.resolveOperand(cond, boolType, result.operands))
+        return failure();
+
+    // Parse the 'then' region.
+    auto parseThenLoc = parser.getCurrentLocation();
+    if (parser.parseRegion(*thenRegion, /*arguments=*/{},
+                           /*argTypes=*/{}))
+        return failure();
+    if (ensureRegionTerm(parser, *thenRegion, parseThenLoc).failed()) return failure();
+
+    // If we find an 'else' keyword, parse the 'else' region.
+    if (!parser.parseOptionalKeyword("else")) {
+        auto parseElseLoc = parser.getCurrentLocation();
+        if (parser.parseRegion(*elseRegion, /*arguments=*/{}, /*argTypes=*/{})) return failure();
+        if (ensureRegionTerm(parser, *elseRegion, parseElseLoc).failed()) return failure();
+    }
+
+    // Parse the optional attribute list.
+    return parser.parseOptionalAttrDict(result.attributes) ? failure() : success();
+}
+
+void P4HIR::IfOp::print(OpAsmPrinter &p) {
+    p << " " << getCondition() << " ";
+    auto &thenRegion = this->getThenRegion();
+    p.printRegion(thenRegion,
+                  /*printEntryBlockArgs=*/false,
+                  /*printBlockTerminators=*/!omitRegionTerm(thenRegion));
+
+    // Print the 'else' regions if it exists and has a block.
+    auto &elseRegion = this->getElseRegion();
+    if (!elseRegion.empty()) {
+        p << " else ";
+        p.printRegion(elseRegion,
+                      /*printEntryBlockArgs=*/false,
+                      /*printBlockTerminators=*/!omitRegionTerm(elseRegion));
+    }
+
+    p.printOptionalAttrDict(getOperation()->getAttrs());
+}
+
+/// Default callback for IfOp builders.
+void P4HIR::buildTerminatedBody(OpBuilder &builder, Location loc) {
+    // add p4hir.yield to the end of the block
+    builder.create<P4HIR::YieldOp>(loc);
+}
+
+void P4HIR::IfOp::getSuccessorRegions(mlir::RegionBranchPoint point,
+                                      SmallVectorImpl<RegionSuccessor> &regions) {
+    // The `then` and the `else` region branch back to the parent operation.
+    if (!point.isParent()) {
+        regions.push_back(RegionSuccessor());
+        return;
+    }
+
+    // Don't consider the else region if it is empty.
+    Region *elseRegion = &this->getElseRegion();
+    if (elseRegion->empty()) elseRegion = nullptr;
+
+    // If the condition isn't constant, both regions may be executed.
+    regions.push_back(RegionSuccessor(&getThenRegion()));
+    // If the else region does not exist, it is not a viable successor.
+    if (elseRegion) regions.push_back(RegionSuccessor(elseRegion));
+}
+
+void P4HIR::IfOp::build(OpBuilder &builder, OperationState &result, Value cond, bool withElseRegion,
+                        function_ref<void(OpBuilder &, Location)> thenBuilder,
+                        function_ref<void(OpBuilder &, Location)> elseBuilder) {
+    assert(thenBuilder && "the builder callback for 'then' must be present");
+
+    result.addOperands(cond);
+
+    OpBuilder::InsertionGuard guard(builder);
+    Region *thenRegion = result.addRegion();
+    builder.createBlock(thenRegion);
+    thenBuilder(builder, result.location);
+
+    Region *elseRegion = result.addRegion();
+    if (!withElseRegion) return;
+
+    builder.createBlock(elseRegion);
+    elseBuilder(builder, result.location);
+}
+
 void P4HIR::P4HIRDialect::initialize() {
     registerTypes();
     registerAttributes();

test/Dialect/P4HIR/if.mlir

@@ -0,0 +1,17 @@
+// RUN: p4mlir-opt %s -o %t.mlir
+// RUN: FileCheck --input-file=%t.mlir %s
+
+module {
+  // No need to check stuff. If it parses, it's fine.
+  // CHECK: module
+  %0 = p4hir.const #p4hir.bool<false> : !p4hir.bool
+  p4hir.if %0 {
+    %29 = p4hir.const #p4hir.bool<true> : !p4hir.bool
+  }
+
+  p4hir.if %0 {
+    %29 = p4hir.const #p4hir.bool<true> : !p4hir.bool
+  } else {
+    %29 = p4hir.const #p4hir.bool<true> : !p4hir.bool
+  }
+}

test/Translate/Ops/if.p4

@@ -0,0 +1,51 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+action iftest() {
+   // CHECK-LABEL: module
+   // CHECK: %[[B:.*]] = p4hir.alloca !p4hir.bool ["b"] : !p4hir.ref<!p4hir.bool>
+   // CHECK: %[[B_VAL:.*]] = p4hir.load %[[B]] : !p4hir.ref<!p4hir.bool>, !p4hir.bool
+   // CHECK: p4hir.if %[[B_VAL]] {
+   // CHECK: }
+   // CHECK: %[[B_VAL:.*]] = p4hir.load %0 : !p4hir.ref<!p4hir.bool>, !p4hir.bool
+   // CHECK: %[[NOT_B_VAL:.*]] = p4hir.unary(not, %[[B_VAL]]) : !p4hir.bool
+   // CHECK: p4hir.if %[[NOT_B_VAL]] {
+   // CHECK: } else {
+   // CHECK: }
+    bool b;
+    if (b) {
+
+    }
+
+    if (!b) {
+
+    } else {
+    }
+}
+
+action iftest2() {
+    int<16> x1 = 1;
+    if (x1 == 2) {
+        int<16> x2 = 3;
+    } else {
+        int<16> x3 = 4;
+    }
+    int<16> x4 = 5;
+}
+
+action iftest3() {
+    int<16> x1 = 1;
+    int<16> x2 = 3;
+    int<16> x3 = 4;
+    if (x1 == 2) {
+        x1 = 2;
+        x2 = 4;
+    } else if (x2 == x3) {
+        x3 = 5;
+        x1 = 7;
+    }
+    if (x1 == x2) {
+        int<16> x4 = x3;
+        x1 = x3;
+    }
+    x2 = 1;
+}

tools/p4mlir-translate/translate.cpp

@@ -340,6 +340,7 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
     bool preorder(const P4::IR::AssignmentStatement *assign) 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;
 
     mlir::Value emitUnOp(const P4::IR::Operation_Unary *unop, P4HIR::UnaryOpKind kind);
     mlir::Value emitBinOp(const P4::IR::Operation_Binary *binop, P4HIR::BinOpKind kind);
@@ -622,6 +623,24 @@ bool P4HIRConverter::preorder(const P4::IR::LAnd *land) {
     return false;
 }
 
+bool P4HIRConverter::preorder(const P4::IR::IfStatement *ifs) {
+    // Materialize condition first
+    visit(ifs->condition);
+
+    // Create if itself
+    builder.create<P4HIR::IfOp>(
+        getLoc(builder, ifs), getValue(ifs->condition), ifs->ifFalse,
+        [&](mlir::OpBuilder &b, mlir::Location) {
+            visit(ifs->ifTrue);
+            P4HIR::buildTerminatedBody(b, getEndLoc(builder, ifs->ifTrue));
+        },
+        [&](mlir::OpBuilder &b, mlir::Location) {
+            visit(ifs->ifFalse);
+            P4HIR::buildTerminatedBody(b, getEndLoc(builder, ifs->ifFalse));
+        });
+    return false;
+}
+
 }  // namespace
 
 namespace P4::P4MLIR {