Implement packet_in.extract logic for hs.next.

lib/Conversion/P4HIRToCoreLib/P4HIRToCoreLib.cpp

@@ -89,6 +89,116 @@ struct CallOpConversionPattern : public OpConversionPattern<P4HIR::CallOp> {
     }
 };
 
+static Value detectNextPattern(P4HIR::CallMethodOp extractOp) {
+    // Pattern:
+    // %hs = ... (HeaderStackType)
+    // %nextIndex_ref = struct_field_ref %hs["nextIndex"]
+    // %idx = read %nextIndex_ref
+    // %elt_ref = array_element_ref %array[%idx]
+    // %var = variable
+    // call_method @extract(%var)      <- extractOp
+    // %val = read %var                <- read from extract result
+    // assign %val, %elt_ref           <- assign extract result to hs[nextIndex]
+
+    if (extractOp.getArgOperands().empty()) return {};
+
+    auto extractArg = extractOp.getArgOperands()[0];
+
+    auto varOp = extractArg.getDefiningOp<P4HIR::VariableOp>();
+    if (!varOp) return {};
+
+    P4HIR::ReadOp readOp = nullptr;
+    for (auto *user : varOp.getResult().getUsers()) {
+        if (auto read = dyn_cast<P4HIR::ReadOp>(user)) {
+            if (read->isBeforeInBlock(extractOp)) continue;
+            readOp = read;
+            break;
+        }
+    }
+    if (!readOp) return {};
+
+    P4HIR::AssignOp assignOp = nullptr;
+    for (auto *user : readOp.getResult().getUsers()) {
+        if (auto assign = dyn_cast<P4HIR::AssignOp>(user)) {
+            assignOp = assign;
+            break;
+        }
+    }
+    if (!assignOp) return {};
+
+    auto target = assignOp.getRef();
+    auto arrayEltRef = target.getDefiningOp<P4HIR::ArrayElementRefOp>();
+    if (!arrayEltRef) return {};
+
+    auto index = arrayEltRef.getIndex();
+    auto readIdx = index.getDefiningOp<P4HIR::ReadOp>();
+    if (!readIdx) return {};
+
+    auto readSource = readIdx.getRef();
+    auto nextIndexRef = readSource.getDefiningOp<P4HIR::StructFieldRefOp>();
+    if (!nextIndexRef || nextIndexRef.getFieldName() != "nextIndex") return {};
+
+    auto hsBase = nextIndexRef.getOperand();
+    auto hsBaseType = hsBase.getType();
+    if (auto refType = mlir::dyn_cast<P4HIR::ReferenceType>(hsBaseType))
+        hsBaseType = refType.getObjectType();
+
+    if (!mlir::isa<P4HIR::HeaderStackType>(hsBaseType)) return {};
+
+    return hsBase;
+}
+
+static void insertHeaderStackExtractLogic(P4HIR::CallMethodOp op, Value hsBase,
+                                          ConversionPatternRewriter &rewriter) {
+    auto loc = op.getLoc();
+    static constexpr unsigned hsIdBitWidth = 32;
+    auto intType = P4HIR::BitsType::get(op.getContext(), hsIdBitWidth, false);
+
+    auto hsBaseType = hsBase.getType();
+    if (auto refType = mlir::dyn_cast<P4HIR::ReferenceType>(hsBaseType))
+        hsBaseType = refType.getObjectType();
+
+    auto hsType = mlir::cast<P4HIR::HeaderStackType>(hsBaseType);
+    uint64_t arraySize = hsType.getArraySize();
+
+    auto nextIndexRef = rewriter.create<P4HIR::StructFieldRefOp>(loc, hsBase, "nextIndex");
+    auto nextIndexVal = rewriter.create<P4HIR::ReadOp>(loc, nextIndexRef);
+    auto sizeConst = rewriter.create<P4HIR::ConstOp>(loc, P4HIR::IntAttr::get(intType, arraySize));
+    auto check = rewriter.create<P4HIR::CmpOp>(loc, P4HIR::CmpOpKind::Lt, nextIndexVal, sizeConst);
+
+    auto errorType = P4HIR::ErrorType::get(
+        op.getContext(),
+        mlir::ArrayAttr::get(op.getContext(),
+                             {mlir::StringAttr::get(op.getContext(), "StackOutOfBounds")}));
+    auto errorConst = rewriter.create<P4HIR::ConstOp>(
+        loc, P4HIR::ErrorCodeAttr::get(errorType, "StackOutOfBounds"));
+    rewriter.create<P4CoreLib::VerifyOp>(loc, check, errorConst);
+
+    auto varOp = op.getArgOperands()[0].getDefiningOp<P4HIR::VariableOp>();
+    if (!varOp) return;
+
+    for (auto *user : varOp.getResult().getUsers()) {
+        if (auto readOp = dyn_cast<P4HIR::ReadOp>(user)) {
+            if (readOp->isBeforeInBlock(op)) continue;
+            for (auto *readUser : readOp.getResult().getUsers()) {
+                if (auto assignOp = dyn_cast<P4HIR::AssignOp>(readUser)) {
+                    rewriter.setInsertionPointAfter(assignOp);
+                    auto constOne =
+                        rewriter.create<P4HIR::ConstOp>(loc, P4HIR::IntAttr::get(intType, 1));
+                    auto nextIndexRefInc =
+                        rewriter.create<P4HIR::StructFieldRefOp>(loc, hsBase, "nextIndex");
+                    auto currentNext = rewriter.create<P4HIR::ReadOp>(loc, nextIndexRefInc);
+                    auto incremented = rewriter.create<P4HIR::BinOp>(loc, P4HIR::BinOpKind::Add,
+                                                                     currentNext, constOne);
+                    rewriter.create<P4HIR::AssignOp>(loc, incremented, nextIndexRefInc);
+                    return;
+                }
+            }
+            break;
+        }
+    }
+}
+
 struct CallMethodOpConversionPattern : public OpConversionPattern<P4HIR::CallMethodOp> {
     using OpConversionPattern<P4HIR::CallMethodOp>::OpConversionPattern;
 
@@ -103,6 +213,11 @@ struct CallMethodOpConversionPattern : public OpConversionPattern<P4HIR::CallMet
         } else if (extSym == "packet_in") {
             if (methodSym == "extract") {
                 size_t sz = op.getArgOperands().size();
+                if (Value hsBase = detectNextPattern(op)) {
+                    auto savedInsertionPoint = rewriter.saveInsertionPoint();
+                    insertHeaderStackExtractLogic(op, hsBase, rewriter);
+                    rewriter.restoreInsertionPoint(savedInsertionPoint);
+                }
                 if (sz == 1) {
                     rewriter.replaceOpWithNewOp<P4CoreLib::PacketExtractOp>(op, op.getResultTypes(),
                                                                             operands.getOperands());

test/Conversion/P4CoreLib/packet_in_hs.mlir

@@ -0,0 +1,53 @@
+// RUN: p4mlir-opt %s --lower-to-p4corelib | FileCheck %s
+!b32i = !p4hir.bit<32>
+!b8i = !p4hir.bit<8>
+!error = !p4hir.error<NoError, StackOutOfBounds>
+!packet_in = !p4hir.extern<"packet_in" annotations {corelib}>
+!validity_bit = !p4hir.validity.bit
+!H = !p4hir.header<"H", data: !b8i, __valid: !validity_bit>
+!arr_2xH = !p4hir.array<2x!H>
+!hs_2xH = !p4hir.header_stack<2x!H>
+!headers_t = !p4hir.struct<"headers_t", stack: !hs_2xH>
+
+module {
+  p4hir.extern @packet_in annotations {corelib} {
+    p4hir.func @extract<!p4hir.type_var<"T">>(!p4hir.ref<!p4hir.type_var<"T">>)
+  }
+
+  // CHECK-LABEL: p4hir.parser @p
+  p4hir.parser @p(%arg0: !packet_in, %arg1: !p4hir.ref<!headers_t>)() {
+    p4hir.state @start {
+      p4hir.scope {
+        %stack_ref = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
+        %data_ref = p4hir.struct_field_ref %stack_ref["data"] : <!hs_2xH>
+        %nextIndex_ref = p4hir.struct_field_ref %stack_ref["nextIndex"] : <!hs_2xH>
+        %nextIndex = p4hir.read %nextIndex_ref : <!b32i>
+        %elt_ref = p4hir.array_element_ref %data_ref[%nextIndex] : !p4hir.ref<!arr_2xH>, !b32i
+        %hdr = p4hir.variable ["hdr"] : <!H>
+
+        // CHECK: p4hir.struct_field_ref {{.*}}["nextIndex"] : <!hs_2xH>
+        // CHECK: p4hir.read {{.*}} : <!b32i>
+        // CHECK: %[[size:.*]] = p4hir.const #int2_b32i
+        // CHECK: %[[check:.*]] = p4hir.cmp(lt, {{.*}} : !b32i, %[[size]] : !b32i)
+        // CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
+        // CHECK: p4corelib.verify %[[check]] signalling %[[error]] : !error
+        // CHECK: p4corelib.extract_header {{.*}} : <!H> from %arg0 : !p4corelib.packet_in
+        p4hir.call_method @packet_in::@extract<[!H]> (%hdr) of %arg0 : !packet_in : (!p4hir.ref<!H>) -> ()
+
+        %val = p4hir.read %hdr : <!H>
+        p4hir.assign %val, %elt_ref : <!H>
+
+        // CHECK: %[[one:.*]] = p4hir.const #int1_b32i
+        // CHECK: p4hir.struct_field_ref {{.*}}["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[curr:.*]] = p4hir.read {{.*}} : <!b32i>
+        // CHECK: %[[inc:.*]] = p4hir.binop(add, %[[curr]], %[[one]]) : !b32i
+        // CHECK: p4hir.assign %[[inc]], {{.*}} : <!b32i>
+      }
+      p4hir.transition to @p::@accept
+    }
+    p4hir.state @accept {
+      p4hir.parser_accept
+    }
+    p4hir.transition to @p::@start
+  }
+}

test/Conversion/P4CoreLib/packet_in_hs.p4

@@ -0,0 +1,47 @@
+// RUN: p4mlir-translate --typeinference-only %s | sed 's/__corelib = \[\]/corelib/g' | p4mlir-opt --lower-to-p4corelib | FileCheck %s
+
+@__corelib
+extern packet_in {
+    void extract<T>(out T hdr);
+}
+
+header H {
+    bit<8> data;
+}
+
+struct headers_t {
+    H[2] stack;
+}
+
+// CHECK-LABEL: p4hir.parser @p
+parser p(packet_in pkt, out headers_t hdr) {
+    // CHECK: p4hir.state @start
+    state start {
+        // CHECK: p4hir.scope
+        // CHECK: %[[stack_field_ref:.*]] = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
+        // CHECK: %[[data_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["data"] : <!hs_2xH>
+        // CHECK: %[[nextIndex_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val:.*]] = p4hir.read %[[nextIndex_field_ref]] : <!b32i>
+        // CHECK: %[[elt_ref:.*]] = p4hir.array_element_ref %[[data_field_ref]][%[[val]]] : !p4hir.ref<!arr_2xH>, !b32i
+        // CHECK: %[[hdr_out_arg:.*]] = p4hir.variable ["hdr_out_arg"] : <!H>
+
+        // CHECK: %[[nextIndex_field_ref_0:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val_1:.*]] = p4hir.read %[[nextIndex_field_ref_0]] : <!b32i>
+        // CHECK: %[[c2_b32i:.*]] = p4hir.const #int2_b32i
+        // CHECK: %[[lt:.*]] = p4hir.cmp(lt, %[[val_1]] : !b32i, %[[c2_b32i]] : !b32i)
+        // CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
+        // CHECK: p4corelib.verify %[[lt]] signalling %[[error]] : !error
+
+        // CHECK: p4corelib.extract_header %[[hdr_out_arg]] : <!H> from %arg0 : !p4corelib.packet_in
+        // CHECK: %[[val_2:.*]] = p4hir.read %[[hdr_out_arg]] : <!H>
+        // CHECK: p4hir.assign %[[val_2]], %[[elt_ref]] : <!H>
+
+        // CHECK: %[[c1_b32i:.*]] = p4hir.const #int1_b32i
+        // CHECK: %[[nextIndex_field_ref_3:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val_4:.*]] = p4hir.read %[[nextIndex_field_ref_3]] : <!b32i>
+        // CHECK: %[[add:.*]] = p4hir.binop(add, %[[val_4]], %[[c1_b32i]]) : !b32i
+        // CHECK: p4hir.assign %[[add]], %[[nextIndex_field_ref_3]] : <!b32i>
+        pkt.extract(hdr.stack.next);
+        transition accept;
+    }
+}

test/Conversion/P4CoreLib/packet_in_hs_varbit.p4

@@ -0,0 +1,49 @@
+// RUN: p4mlir-translate --typeinference-only %s | sed 's/__corelib = \[\]/corelib/g' | p4mlir-opt --lower-to-p4corelib | FileCheck %s
+
+@__corelib
+extern packet_in {
+    void extract<T>(out T hdr, in bit<32> variableFieldSizeInBits);
+}
+
+header H {
+    bit<8> fixed_data;
+    varbit<64> variable_data;
+}
+
+struct headers_t {
+    H[2] stack;
+}
+
+// CHECK-LABEL: p4hir.parser @p
+parser p(packet_in pkt, out headers_t hdr) {
+    // CHECK: p4hir.state @start
+    state start {
+        // CHECK: p4hir.scope
+        // CHECK: %[[stack_field_ref:.*]] = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
+        // CHECK: %[[data_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["data"] : <!hs_2xH>
+        // CHECK: %[[nextIndex_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val:.*]] = p4hir.read %[[nextIndex_field_ref]] : <!b32i>
+        // CHECK: %[[elt_ref:.*]] = p4hir.array_element_ref %[[data_field_ref]][%[[val]]] : !p4hir.ref<!arr_2xH>, !b32i
+        // CHECK: %[[hdr_out_arg:.*]] = p4hir.variable ["hdr_out_arg"] : <!H>
+
+        // Bounds check inserted by lowering
+        // CHECK: %[[nextIndex_field_ref_0:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val_1:.*]] = p4hir.read %[[nextIndex_field_ref_0]] : <!b32i>
+        // CHECK: %[[c2_b32i:.*]] = p4hir.const #int2_b32i
+        // CHECK: %[[lt:.*]] = p4hir.cmp(lt, %[[val_1]] : !b32i, %[[c2_b32i]] : !b32i)
+        // CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
+        // CHECK: p4corelib.verify %[[lt]] signalling %[[error]] : !error
+
+        // Variable-width extract lowered to p4corelib
+        // CHECK: p4corelib.extract_header_variable %[[hdr_out_arg]]<{{.*}}> : <!H> from %arg0 : !p4corelib.packet_in
+
+        // nextIndex increment inserted by lowering
+        // CHECK: %[[c1_b32i:.*]] = p4hir.const #int1_b32i
+        // CHECK: %[[nextIndex_field_ref_3:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
+        // CHECK: %[[val_4:.*]] = p4hir.read %[[nextIndex_field_ref_3]] : <!b32i>
+        // CHECK: %[[add:.*]] = p4hir.binop(add, %[[val_4]], %[[c1_b32i]]) : !b32i
+        // CHECK: p4hir.assign %[[add]], %[[nextIndex_field_ref_3]] : <!b32i>
+        pkt.extract(hdr.stack.next, 32);
+        transition accept;
+    }
+}