Add generic scope op. Lower BlockStatement to it

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.h

@@ -9,6 +9,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -3,6 +3,7 @@
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -324,4 +325,121 @@ def CmpOp : P4HIR_Op<"cmp", [Pure, SameTypeOperands]> {
   let hasVerifier = 0;
 }
 
+def ScopeOp : P4HIR_Op<"scope", [
+       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+       RecursivelySpeculatable, AutomaticAllocationScope,
+       NoRegionArguments]> {
+  let summary = "Represents a P4 scope";
+  let description = [{
+    `p4hir.scope` contains one region and defines a strict "scope" for all new
+    values produced within its blocks.
+
+    The region can contain an arbitrary number of blocks but usually defaults
+    to one and can optionally return a value (useful for representing values
+    coming out of expressions that are not parts of other expressions, short-cut
+    logical operations and simplifying inlining logic) via `p4hir.yield`:
+
+
+    ```mlir
+    %rvalue = p4hir.scope {
+      ...
+      p4hir.yield %value
+    }
+    ```
+
+    The blocks can be terminated by `p4hir.yield` or `p4hir.return`.
+    If `p4hir.scope` yields no value, the `p4hir.yield` can be left out, and
+    will be inserted implicitly.
+  }];
+
+  let results = (outs Optional<AnyP4Type>:$results);
+  let regions = (region AnyRegion:$scopeRegion);
+
+  let hasVerifier = 1;
+  let skipDefaultBuilders = 1;
+  let assemblyFormat = [{
+    custom<OmittedTerminatorRegion>($scopeRegion) (`:` type($results)^)? attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// Determine whether the scope is empty, meaning it contains a single block
+    /// terminated by a p4hir.yield.
+    bool isEmpty() {
+      auto &entry = getScopeRegion().front();
+      return getScopeRegion().hasOneBlock() &&
+        llvm::isa<YieldOp>(entry.front());
+      }
+    }];
+
+  let builders = [
+    // Scopes for yielding values.
+    OpBuilder<(ins
+              "llvm::function_ref<void(mlir::OpBuilder &, mlir::Type &, mlir::Location)>":$scopeBuilder)>,
+    // Scopes without yielding values.
+    OpBuilder<(ins "llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)>":$scopeBuilder)>
+  ];
+}
+
+def YieldOp : P4HIR_Op<"yield", [ReturnLike, Terminator,
+    ParentOneOf<["ScopeOp",
+                 // "IfOp", "TernaryOp"
+                 // "SwitchOp", "CaseOp",
+                 // "ForInOp", "ForOp",
+                 // "CallOp"
+                 ]>]> {
+  let summary = "Represents the default branching behaviour of a region";
+  let description = [{
+    The `p4hir.yield` operation terminates regions on different P4HIR operations,
+    and it is used to represent the default branching behaviour of a region.
+    Said branching behaviour is determined by the parent operation. For
+    example, a yield in a `p4hir.if` region implies a branch to the exit block,
+    and so on.
+
+    In some cases, it might yield an SSA value and the semantics of how the
+    values are yielded is defined by the parent operation. For example, a
+    `p4hir.ternary` operation yields a value from one of its regions.
+
+    As a general rule, `p4hir.yield` must be explicitly used whenever a region has
+    more than one block and no terminator, or within `p4hir.switch` regions not
+    `p4hir.return` terminated.
+
+    Examples:
+    ```mlir
+    p4hir.if %4 {
+      ...
+      p4hir.yield
+    }
+
+    p4hir.switch (%5) [
+      case (equal, 3) {
+        ...
+        p4hir.yield
+      }, ...
+    ]
+
+    p4hir.scope {
+      ...
+      p4hir.yield
+    }
+
+    %x = p4hir.scope {
+      ...
+      p4hir.yield %val
+    }
+
+    %y = p4hir.ternary {
+      ...
+      p4hir.yield %val : !p4hir.bit<42>
+    } :  !p4hir.bit<42>
+    ```
+  }];
+
+  let arguments = (ins Variadic<AnyP4Type>:$args);
+  let assemblyFormat = "($args^ `:` type($args))? attr-dict";
+  let builders = [
+    OpBuilder<(ins), [{ /* nothing to do */ }]>,
+  ];
+}
+
+
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_OPS_TD

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -7,11 +7,6 @@
 #include "p4mlir/Dialect/P4HIR/P4HIR_OpsEnums.h"
 #include "p4mlir/Dialect/P4HIR/P4HIR_Types.h"
 
-#define GET_OP_CLASSES
-#include "p4mlir/Dialect/P4HIR/P4HIR_Dialect.cpp.inc"
-#include "p4mlir/Dialect/P4HIR/P4HIR_Ops.cpp.inc"
-#include "p4mlir/Dialect/P4HIR/P4HIR_OpsEnums.cpp.inc"
-
 using namespace mlir;
 using namespace P4::P4MLIR;
 
@@ -73,6 +68,110 @@ void P4HIR::AllocaOp::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationStat
     odsState.addTypes(ref);
 }
 
+//===----------------------------------------------------------------------===//
+// ScopeOp
+//===----------------------------------------------------------------------===//
+
+void P4HIR::ScopeOp::getSuccessorRegions(mlir::RegionBranchPoint point,
+                                         SmallVectorImpl<RegionSuccessor> &regions) {
+    // The only region always branch back to the parent operation.
+    if (!point.isParent()) {
+        regions.push_back(RegionSuccessor(getODSResults(0)));
+        return;
+    }
+
+    // If the condition isn't constant, both regions may be executed.
+    regions.push_back(RegionSuccessor(&getScopeRegion()));
+}
+
+void P4HIR::ScopeOp::build(OpBuilder &builder, OperationState &result,
+                           function_ref<void(OpBuilder &, Type &, Location)> scopeBuilder) {
+    assert(scopeBuilder && "the builder callback for 'then' must be present");
+
+    OpBuilder::InsertionGuard guard(builder);
+    Region *scopeRegion = result.addRegion();
+    builder.createBlock(scopeRegion);
+
+    mlir::Type yieldTy;
+    scopeBuilder(builder, yieldTy, result.location);
+
+    if (yieldTy) result.addTypes(TypeRange{yieldTy});
+}
+
+void P4HIR::ScopeOp::build(OpBuilder &builder, OperationState &result,
+                           function_ref<void(OpBuilder &, Location)> scopeBuilder) {
+    assert(scopeBuilder && "the builder callback for 'then' must be present");
+    OpBuilder::InsertionGuard guard(builder);
+    Region *scopeRegion = result.addRegion();
+    builder.createBlock(scopeRegion);
+    scopeBuilder(builder, result.location);
+}
+
+LogicalResult P4HIR::ScopeOp::verify() {
+    if (getScopeRegion().empty()) {
+        return emitOpError() << "p4hir.scope must not be empty since it should "
+                                "include at least an implicit p4hir.yield ";
+    }
+
+    if (getScopeRegion().back().empty() || !getScopeRegion().back().mightHaveTerminator() ||
+        !getScopeRegion().back().getTerminator()->hasTrait<OpTrait::IsTerminator>())
+        return emitOpError() << "last block of p4hir.scope must be terminated";
+    return success();
+}
+//===----------------------------------------------------------------------===//
+// Custom Parsers & Printers
+//===----------------------------------------------------------------------===//
+
+// Check if a region's termination omission is valid and, if so, creates and
+// inserts the omitted terminator into the region.
+static LogicalResult ensureRegionTerm(OpAsmParser &parser, Region &region, SMLoc errLoc) {
+    Location eLoc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
+    OpBuilder builder(parser.getBuilder().getContext());
+
+    // Insert empty block in case the region is empty to ensure the terminator
+    // will be inserted
+    if (region.empty()) builder.createBlock(&region);
+
+    Block &block = region.back();
+    // Region is properly terminated: nothing to do.
+    if (!block.empty() && block.back().hasTrait<OpTrait::IsTerminator>()) return success();
+
+    // Check for invalid terminator omissions.
+    if (!region.hasOneBlock())
+        return parser.emitError(errLoc, "multi-block region must not omit terminator");
+
+    // Terminator was omitted correctly: recreate it.
+    builder.setInsertionPointToEnd(&block);
+    builder.create<P4HIR::YieldOp>(eLoc);
+    return success();
+}
+
+static mlir::ParseResult parseOmittedTerminatorRegion(mlir::OpAsmParser &parser,
+                                                      mlir::Region &scopeRegion) {
+    auto regionLoc = parser.getCurrentLocation();
+    if (parser.parseRegion(scopeRegion)) return failure();
+    if (ensureRegionTerm(parser, scopeRegion, regionLoc).failed()) return failure();
+
+    return success();
+}
+
+// True if the region's terminator should be omitted.
+bool omitRegionTerm(mlir::Region &r) {
+    const auto singleNonEmptyBlock = r.hasOneBlock() && !r.back().empty();
+    const auto yieldsNothing = [&r]() {
+        auto y = dyn_cast<P4HIR::YieldOp>(r.back().getTerminator());
+        return y && y.getArgs().empty();
+    };
+    return singleNonEmptyBlock && yieldsNothing();
+}
+
+static void printOmittedTerminatorRegion(mlir::OpAsmPrinter &printer, P4HIR::ScopeOp &,
+                                         mlir::Region &scopeRegion) {
+    printer.printRegion(scopeRegion,
+                        /*printEntryBlockArgs=*/false,
+                        /*printBlockTerminators=*/!omitRegionTerm(scopeRegion));
+}
+
 void P4HIR::P4HIRDialect::initialize() {
     registerTypes();
     registerAttributes();
@@ -81,3 +180,8 @@ void P4HIR::P4HIRDialect::initialize() {
 #include "p4mlir/Dialect/P4HIR/P4HIR_Ops.cpp.inc"  // NOLINT
         >();
 }
+
+#define GET_OP_CLASSES
+#include "p4mlir/Dialect/P4HIR/P4HIR_Dialect.cpp.inc"
+#include "p4mlir/Dialect/P4HIR/P4HIR_Ops.cpp.inc"
+#include "p4mlir/Dialect/P4HIR/P4HIR_OpsEnums.cpp.inc"

test/Dialect/P4HIR/scope.mlir

@@ -0,0 +1,34 @@
+// RUN: p4mlir-opt %s -o %t.mlir
+// RUN: FileCheck --input-file=%t.mlir %s
+
+!bit32 = !p4hir.bit<32>
+
+module {
+  // Should properly print/parse scope with implicit empty yield.
+  // TODO: When functions will be ready, uncomment lines below
+  // xHECK-LABEL: implicit_yield
+  // p4hir.func @implicit_yield() {
+    p4hir.scope {
+    }
+    // CHECK: p4hir.scope {
+    // CHECK-NEXT: }
+    // xHECK-NEXT: p4hir.return
+    // p4hir.return
+  // }
+
+  // Should properly print/parse scope with explicit yield.
+  // xHECK-LABEL: explicit_yield
+  // p4hir.func @explicit_yield() {
+    %0 = p4hir.scope {
+      %1 = p4hir.alloca !bit32 ["a", init] : !p4hir.ref<!bit32>
+      %2 = p4hir.load %1 : !p4hir.ref<!bit32>, !bit32
+      p4hir.yield %2 : !bit32
+    } : !bit32
+    // CHECK: %0 = p4hir.scope {
+    //          [...]
+    // CHECK:   p4hir.yield %2 : !p4hir.bit<32>
+    // CHECK: } : !p4hir.bit<32>
+  //  p4hir.return
+  //}
+
+}

test/Translate/Ops/scope.p4

@@ -0,0 +1,30 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+// CHECK-LABEL: module
+action scope() {
+    bool res;
+    // Outer alloca
+    // CHECK-NEXT: %[[RES_0:.*]] = p4hir.alloca !p4hir.bool ["res"] : !p4hir.ref<!p4hir.bool>
+
+    {
+      bit<10> lhs = 1;
+      bit<10> rhs = 2;
+
+      res = lhs == rhs;
+      // CHECK: p4hir.store %{{.*}}, %[[RES_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+      {
+         res = lhs != rhs;
+      // CHECK: p4hir.store %{{.*}}, %[[RES_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+      }
+    }
+
+    {
+      bit<10> lhs = 1;
+      bit<10> rhs = 2;
+      // CHECK: %[[RES_1:.*]] = p4hir.alloca !p4hir.bit<10> ["res"] : !p4hir.ref<!p4hir.bit<10>>      
+      bit<10> res;
+      // This should store into inner res, not outer
+      // CHECK: p4hir.store %{{.*}}, %[[RES_1]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+      res = lhs * rhs;
+    }
+}

tools/p4mlir-translate/translate.cpp

@@ -47,6 +47,18 @@ mlir::Location getLoc(mlir::OpBuilder &builder, const P4::IR::Node *node) {
         start.getColumnNumber());
 }
 
+mlir::Location getEndLoc(mlir::OpBuilder &builder, const P4::IR::Node *node) {
+    CHECK_NULL(node);
+    auto sourceInfo = node->getSourceInfo();
+    if (!sourceInfo.isValid()) return mlir::UnknownLoc::get(builder.getContext());
+
+    const auto &end = sourceInfo.getEnd();
+
+    return mlir::FileLineColLoc::get(
+        builder.getStringAttr(sourceInfo.getSourceFile().string_view()), end.getLineNumber(),
+        end.getColumnNumber());
+}
+
 mlir::APInt toAPInt(const P4::big_int &value, unsigned bitWidth = 0) {
     std::vector<uint64_t> valueBits;
     // Export absolute value into 64-bit unsigned values, most significant bit last
@@ -247,13 +259,26 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
 
     bool preorder(const P4::IR::P4Program *) override { return true; }
     bool preorder(const P4::IR::P4Action *a) override {
-        // TODO: For now simply visit every node of the body
+        // We cannot simply visit each node of the top-level block as
+        // ResolutionContext would not be able to resolve declarations there
+        // (sic!)
         visit(a->body);
         return false;
     }
     bool preorder(const P4::IR::BlockStatement *block) override {
-        // TODO: For now simply visit every node of the block, create scope afterwards
-        visit(block->components);
+        // If this is a top-level block where scope is implied (e.g. function,
+        // action, certain statements) do not create explicit scope.
+        if (getParent<P4::IR::BlockStatement>()) {
+            mlir::OpBuilder::InsertionGuard guard(builder);
+            auto scope = builder.create<P4HIR::ScopeOp>(
+                getLoc(builder, block),                   /*scopeBuilder=*/
+                [&](mlir::OpBuilder &, mlir::Location) {  // nothing is being yielded
+                    visit(block->components);
+                });
+            builder.setInsertionPointToEnd(&scope.getScopeRegion().back());
+            builder.create<P4HIR::YieldOp>(getEndLoc(builder, block));
+        } else
+            visit(block->components);
         return false;
     }