Add possibility to instantiate subparsers

Signed-off-by: Anton Korobeynikov <anton@korobeynikov.info>

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -996,6 +996,65 @@ def TupleExtractOp : P4HIR_Op<"tuple_extract",
   let hasVerifier = 1;
 }
 
+def InstantiateOp : P4HIR_Op<"instantiate",
+  [NoRegionArguments, CallOpInterface,
+   DeclareOpInterfaceMethods<SymbolUserOpInterface>,
+   DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+  let summary = "constructor call operation";
+  let description = [{
+    Instantiate P4 object: extern, parser, control or package/
+
+    Example:
+
+    ```mlir
+    // Direct call of function
+    %2 = p4hir.call @my_add(%0, %1) : (!p4hir.bit<8>, !p4hir.bit<8>) -> !p4hir.bit<8>
+    ...
+    ```
+  }];
+
+  // TODO: Refine result types, refine parameter type (must be constants!)
+  let results = (outs P4ObjectType:$result);
+  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands,
+                       OptionalAttr<StrAttr>:$name);
+
+  // TODO: Add verifier
+  let hasVerifier = 0;
+
+  let extraClassDeclaration = [{
+    /// Get the argument operands to the called function.
+    mlir::OperandRange getArgOperands() {
+      return {arg_operand_begin(), arg_operand_end()};
+    }
+
+    mlir::MutableOperandRange getArgOperandsMutable() {
+      return getOperandsMutable();
+    }
+
+    operand_iterator arg_operand_begin() { return operand_begin(); }
+    operand_iterator arg_operand_end() { return operand_end(); }
+
+    /// Return the callee of this operation.
+    mlir::CallInterfaceCallable getCallableForCallee() {
+      return (*this)->getAttrOfType<mlir::SymbolRefAttr>("callee");
+    }
+
+    /// Set the callee for this operation.
+    void setCalleeFromCallable(mlir::CallInterfaceCallable callee) {
+      (*this)->setAttr("callee", callee.get<mlir::SymbolRefAttr>());
+    }
+
+    void setArg(unsigned index, mlir::Value value) {
+      setOperand(index, value);
+    }
+  }];
+
+  let assemblyFormat = [{
+    $callee `(` $operands `)` `as` $name attr-dict `:` functional-type($operands, results)
+  }];
+}
+
+
 include "p4mlir/Dialect/P4HIR/P4HIR_ParserOps.td"
 
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_OPS_TD

include/p4mlir/Dialect/P4HIR/P4HIR_Types.td

@@ -572,6 +572,7 @@ def LoadableP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType,
                                 ValidBitType]> {}
 def AnyEnumType : AnyTypeOf<[EnumType, SerEnumType]>;
 def StructLikeType : AnyTypeOf<[StructType, HeaderType]>;
+def P4ObjectType : AnyTypeOf<[ParserType]>;
 
 /// A ref type with the specified constraints on the nested type.
 class SpecificRefType<Type type> : ConfinedType<ReferenceType,

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -1296,6 +1296,43 @@ void P4HIR::MaskOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
     setNameFn(getResult(), "mask");
 }
 
+//===----------------------------------------------------------------------===//
+// InstantiateOp
+//===----------------------------------------------------------------------===//
+
+void P4HIR::InstantiateOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
+    if (getName() && !getName()->empty()) setNameFn(getResult(), *getName());
+}
+
+LogicalResult P4HIR::InstantiateOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+    // Check that the callee attribute was specified.
+    auto ctorAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
+    if (!ctorAttr) return emitOpError("requires a 'callee' symbol reference attribute");
+
+    ParserOp parser =
+        symbolTable.lookupNearestSymbolFrom<ParserOp>(getParentModule(*this), ctorAttr);
+    if (!parser)
+        return emitOpError() << "'" << ctorAttr.getValue() << "' does not reference a valid parser";
+
+    // Verify that the operand and result types match the callee.
+    auto ctorType = parser.getCtorType();
+    if (ctorType.getNumInputs() != getNumOperands())
+        return emitOpError("incorrect number of operands for callee");
+
+    for (unsigned i = 0, e = ctorType.getNumInputs(); i != e; ++i)
+        if (getOperand(i).getType() != ctorType.getInput(i))
+            return emitOpError("operand type mismatch: expected operand type ")
+                   << ctorType.getInput(i) << ", but provided " << getOperand(i).getType()
+                   << " for operand number " << i;
+
+    // Parser itself and return value types must match.
+    if (getResult().getType() != ctorType.getReturnType())
+        return emitOpError("result type mismatch: expected ")
+               << ctorType.getReturnType() << ", but provided " << getResult().getType();
+
+    return success();
+}
+
 namespace {
 struct P4HIROpAsmDialectInterface : public OpAsmDialectInterface {
     using OpAsmDialectInterface::OpAsmDialectInterface;

test/Dialect/P4HIR/instantiate.mlir

@@ -0,0 +1,65 @@
+// RUN: p4mlir-opt %s | FileCheck %s
+
+!empty = !p4hir.struct<"empty">
+!i10i = !p4hir.int<10>
+#false = #p4hir.bool<false> : !p4hir.bool
+#subparser2_ctorArg = #p4hir.ctor_param<@subparser2, "ctorArg"> : !p4hir.bool
+#int1_i10i = #p4hir.int<1> : !i10i
+#int2_i10i = #p4hir.int<2> : !i10i
+// CHECK: module
+module {
+  p4hir.parser @subparser(%arg0: !empty)() {
+    p4hir.state @start {
+      p4hir.transition to @subparser::@accept
+    }
+    p4hir.state @accept {
+      p4hir.parser_accept
+    }
+    p4hir.state @reject {
+      p4hir.parser_reject
+    }
+    p4hir.transition to @subparser::@start
+  }
+  p4hir.parser @subparser2(%arg0: !empty)(ctorArg: !p4hir.bool) {
+    %ctorArg = p4hir.const ["ctorArg"] #subparser2_ctorArg
+    p4hir.state @start {
+      p4hir.transition to @subparser2::@accept
+    }
+    p4hir.state @accept {
+      p4hir.parser_accept
+    }
+    p4hir.state @reject {
+      p4hir.parser_reject
+    }
+    p4hir.transition to @subparser2::@start
+  }
+  p4hir.parser @p(%arg0: !empty, %arg1: !i10i)() {
+    %s = p4hir.variable ["s", init] : <!i10i>
+    p4hir.assign %arg1, %s : <!i10i>
+    %sp = p4hir.instantiate @subparser() as "sp" : () -> !p4hir.parser<"subparser", (!empty)>
+    %false = p4hir.const #false
+    %sp2 = p4hir.instantiate @subparser2(%false) as "sp2" : (!p4hir.bool) -> !p4hir.parser<"subparser2", (!empty)>
+    p4hir.state @start {
+      %c1_i10i = p4hir.const #int1_i10i
+      %cast = p4hir.cast(%c1_i10i : !i10i) : !i10i
+      p4hir.assign %cast, %s : <!i10i>
+      p4hir.transition to @p::@next
+    }
+    p4hir.state @next {
+      %c2_i10i = p4hir.const #int2_i10i
+      %cast = p4hir.cast(%c2_i10i : !i10i) : !i10i
+      p4hir.assign %cast, %s : <!i10i>
+      p4hir.transition to @p::@accept
+    }
+    p4hir.state @drop {
+      p4hir.transition to @p::@reject
+    }
+    p4hir.state @accept {
+      p4hir.parser_accept
+    }
+    p4hir.state @reject {
+      p4hir.parser_reject
+    }
+    p4hir.transition to @p::@start
+  }
+}

test/Translate/Parser/subparser.p4

@@ -0,0 +1,38 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+struct empty {}
+
+parser subparser(in empty e) {
+    state start {
+        transition accept;
+    }
+}
+
+parser subparser2(in empty e)(bool ctorArg) {
+    state start {
+        transition accept;
+    }
+}
+
+// CHECK-LABEL: p4hir.parser @p
+// CHECK: p4hir.instantiate @subparser() as "sp" : () -> !p4hir.parser<"subparser", (!empty)>
+// CHECK: %[[false:.*]] = p4hir.const #false
+// CHECK: p4hir.instantiate @subparser2(%[[false]]) as "sp2" : (!p4hir.bool) -> !p4hir.parser<"subparser2", (!empty)>
+
+parser p(in empty e, in int<10> sinit) {
+    int<10> s = sinit;
+    subparser() sp;
+    subparser2(false) sp2;
+
+    state start {
+        s = 1;
+        transition next;
+    }
+
+    state next {
+        s = 2;
+        transition accept;
+    }
+
+    state drop {}
+}

tools/p4mlir-translate/translate.cpp

@@ -3,8 +3,6 @@
 #include <algorithm>
 #include <climits>
 
-#include "ir/ir-generated.h"
-
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wcovered-switch-default"
 #include "frontends/common/resolveReferences/resolveReferences.h"
@@ -433,6 +431,7 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
     void postorder(const P4::IR::Member *m) override;
 
     bool preorder(const P4::IR::Declaration_Constant *decl) override;
+    bool preorder(const P4::IR::Declaration_Instance *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;
@@ -496,6 +495,8 @@ bool P4TypeConverter::preorder(const P4::IR::Type_Name *name) {
     ConversionTracer trace("Resolving type by name ", name);
     const auto *type = converter.resolveType(name);
     CHECK_NULL(type);
+    LOG4("Resolved to: " << dbp(type));
+
     mlir::Type mlirType = convert(type);
     return setType(name, mlirType);
 }
@@ -1676,6 +1677,38 @@ bool P4HIRConverter::preorder(const P4::IR::SelectExpression *select) {
     return false;
 }
 
+bool P4HIRConverter::preorder(const P4::IR::Declaration_Instance *decl) {
+    ConversionTracer trace("Converting ", decl);
+
+    // P4::Instantiation goes via typeMap and it returns some weird clone
+    // instead of converted type
+    const auto *type = resolveType(decl->type)->to<P4::IR::Type_Declaration>();
+    CHECK_NULL(type);
+    LOG4("Resolved to: " << dbp(type));
+
+    llvm::SmallVector<mlir::Value, 4> operands;
+    for (const auto *arg : *decl->arguments) {
+        ConversionTracer trace("Converting ", arg);
+        operands.push_back(materializeConstantExpr(arg->expression));
+    }
+
+    if (const auto *parser = type->to<P4::IR::P4Parser>()) {
+        LOG4("resolved as parser instantiation");
+        auto resultType = getOrCreateType(parser->getConstructorMethodType()->returnType);
+
+        auto parserSym = p4Symbols.lookup(parser);
+        BUG_CHECK(parserSym, "expected reference parser to be converted: %1%", dbp(parser));
+
+        builder.create<P4HIR::InstantiateOp>(getLoc(builder, decl), resultType,
+                                             parserSym.getRootReference(), operands,
+                                             builder.getStringAttr(decl->name.string_view()));
+    } else {
+        BUG("unsupported instance type: %1%", decl);
+    }
+
+    return false;
+}
+
 }  // namespace
 
 namespace P4::P4MLIR {