Skip to content

Commit e7880b3

Browse files
vaalfrejavaalfreja
committed
Implement correct packet_in.extract lowering.
Signed-off-by: Julia Koval <c_yuliak@xsightlabs.com>
1 parent ff53252 commit e7880b3

File tree

4 files changed

+264
-0
lines changed

4 files changed

+264
-0
lines changed

‎lib/Conversion/P4HIRToCoreLib/P4HIRToCoreLib.cpp‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,116 @@ struct CallOpConversionPattern : public OpConversionPattern<P4HIR::CallOp> {
8989
}
9090
};
9191

92+
static std::optional<Value> detectNextPattern(P4HIR::CallMethodOp extractOp) {
93+
// Pattern:
94+
// %hs = ... (HeaderStackType)
95+
// %nextIndex_ref = struct_field_ref %hs["nextIndex"]
96+
// %idx = read %nextIndex_ref
97+
// %elt_ref = array_element_ref %array[%idx]
98+
// %var = variable
99+
// call_method @extract(%var) <- extractOp
100+
// %val = read %var <- read from extract result
101+
// assign %val, %elt_ref <- assign extract result to hs[nextIndex]
102+
103+
if (extractOp.getArgOperands().empty()) return std::nullopt;
104+
105+
auto extractArg = extractOp.getArgOperands()[0];
106+
107+
auto varOp = extractArg.getDefiningOp<P4HIR::VariableOp>();
108+
if (!varOp) return std::nullopt;
109+
110+
P4HIR::ReadOp readOp = nullptr;
111+
for (auto *user : varOp.getResult().getUsers()) {
112+
if (auto read = dyn_cast<P4HIR::ReadOp>(user)) {
113+
if (read->isBeforeInBlock(extractOp)) continue;
114+
readOp = read;
115+
break;
116+
}
117+
}
118+
if (!readOp) return std::nullopt;
119+
120+
P4HIR::AssignOp assignOp = nullptr;
121+
for (auto *user : readOp.getResult().getUsers()) {
122+
if (auto assign = dyn_cast<P4HIR::AssignOp>(user)) {
123+
assignOp = assign;
124+
break;
125+
}
126+
}
127+
if (!assignOp) return std::nullopt;
128+
129+
auto target = assignOp.getRef();
130+
auto arrayEltRef = target.getDefiningOp<P4HIR::ArrayElementRefOp>();
131+
if (!arrayEltRef) return std::nullopt;
132+
133+
auto index = arrayEltRef.getIndex();
134+
auto readIdx = index.getDefiningOp<P4HIR::ReadOp>();
135+
if (!readIdx) return std::nullopt;
136+
137+
auto readSource = readIdx.getRef();
138+
auto nextIndexRef = readSource.getDefiningOp<P4HIR::StructFieldRefOp>();
139+
if (!nextIndexRef || nextIndexRef.getFieldName() != "nextIndex") return std::nullopt;
140+
141+
auto hsBase = nextIndexRef.getOperand();
142+
auto hsBaseType = hsBase.getType();
143+
if (auto refType = mlir::dyn_cast<P4HIR::ReferenceType>(hsBaseType))
144+
hsBaseType = refType.getObjectType();
145+
146+
if (!mlir::isa<P4HIR::HeaderStackType>(hsBaseType)) return std::nullopt;
147+
148+
return hsBase;
149+
}
150+
151+
static void insertHeaderStackExtractLogic(P4HIR::CallMethodOp op, Value hsBase,
152+
ConversionPatternRewriter &rewriter) {
153+
auto loc = op.getLoc();
154+
static constexpr unsigned hsIdBitWidth = 32;
155+
auto intType = P4HIR::BitsType::get(op.getContext(), hsIdBitWidth, false);
156+
157+
auto hsBaseType = hsBase.getType();
158+
if (auto refType = mlir::dyn_cast<P4HIR::ReferenceType>(hsBaseType))
159+
hsBaseType = refType.getObjectType();
160+
161+
auto hsType = mlir::cast<P4HIR::HeaderStackType>(hsBaseType);
162+
uint64_t arraySize = hsType.getArraySize();
163+
164+
auto nextIndexRef = rewriter.create<P4HIR::StructFieldRefOp>(loc, hsBase, "nextIndex");
165+
auto nextIndexVal = rewriter.create<P4HIR::ReadOp>(loc, nextIndexRef);
166+
auto sizeConst = rewriter.create<P4HIR::ConstOp>(loc, P4HIR::IntAttr::get(intType, arraySize));
167+
auto check = rewriter.create<P4HIR::CmpOp>(loc, P4HIR::CmpOpKind::Lt, nextIndexVal, sizeConst);
168+
169+
auto errorType = P4HIR::ErrorType::get(
170+
op.getContext(),
171+
mlir::ArrayAttr::get(op.getContext(),
172+
{mlir::StringAttr::get(op.getContext(), "StackOutOfBounds")}));
173+
auto errorConst = rewriter.create<P4HIR::ConstOp>(
174+
loc, P4HIR::ErrorCodeAttr::get(errorType, "StackOutOfBounds"));
175+
rewriter.create<P4CoreLib::VerifyOp>(loc, check, errorConst);
176+
177+
auto varOp = op.getArgOperands()[0].getDefiningOp<P4HIR::VariableOp>();
178+
if (!varOp) return;
179+
180+
for (auto *user : varOp.getResult().getUsers()) {
181+
if (auto readOp = dyn_cast<P4HIR::ReadOp>(user)) {
182+
if (readOp->isBeforeInBlock(op)) continue;
183+
for (auto *readUser : readOp.getResult().getUsers()) {
184+
if (auto assignOp = dyn_cast<P4HIR::AssignOp>(readUser)) {
185+
rewriter.setInsertionPointAfter(assignOp);
186+
auto constOne =
187+
rewriter.create<P4HIR::ConstOp>(loc, P4HIR::IntAttr::get(intType, 1));
188+
auto nextIndexRefInc =
189+
rewriter.create<P4HIR::StructFieldRefOp>(loc, hsBase, "nextIndex");
190+
auto currentNext = rewriter.create<P4HIR::ReadOp>(loc, nextIndexRefInc);
191+
auto incremented = rewriter.create<P4HIR::BinOp>(loc, P4HIR::BinOpKind::Add,
192+
currentNext, constOne);
193+
rewriter.create<P4HIR::AssignOp>(loc, incremented, nextIndexRefInc);
194+
return;
195+
}
196+
}
197+
break;
198+
}
199+
}
200+
}
201+
92202
struct CallMethodOpConversionPattern : public OpConversionPattern<P4HIR::CallMethodOp> {
93203
using OpConversionPattern<P4HIR::CallMethodOp>::OpConversionPattern;
94204

@@ -103,6 +213,11 @@ struct CallMethodOpConversionPattern : public OpConversionPattern<P4HIR::CallMet
103213
} else if (extSym == "packet_in") {
104214
if (methodSym == "extract") {
105215
size_t sz = op.getArgOperands().size();
216+
if (auto hsBase = detectNextPattern(op)) {
217+
auto savedInsertionPoint = rewriter.saveInsertionPoint();
218+
insertHeaderStackExtractLogic(op, *hsBase, rewriter);
219+
rewriter.restoreInsertionPoint(savedInsertionPoint);
220+
}
106221
if (sz == 1) {
107222
rewriter.replaceOpWithNewOp<P4CoreLib::PacketExtractOp>(op, op.getResultTypes(),
108223
operands.getOperands());

‎test/Conversion/P4CoreLib/packet_in_hs.mlir‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
// RUN: p4mlir-opt %s --lower-to-p4corelib | FileCheck %s
2+
!b32i = !p4hir.bit<32>
3+
!b8i = !p4hir.bit<8>
4+
!error = !p4hir.error<NoError, StackOutOfBounds>
5+
!packet_in = !p4hir.extern<"packet_in" annotations {corelib}>
6+
!validity_bit = !p4hir.validity.bit
7+
!H = !p4hir.header<"H", data: !b8i, __valid: !validity_bit>
8+
!arr_2xH = !p4hir.array<2x!H>
9+
!hs_2xH = !p4hir.header_stack<2x!H>
10+
!headers_t = !p4hir.struct<"headers_t", stack: !hs_2xH>
11+
12+
module {
13+
p4hir.extern @packet_in annotations {corelib} {
14+
p4hir.func @extract<!p4hir.type_var<"T">>(!p4hir.ref<!p4hir.type_var<"T">>)
15+
}
16+
17+
// CHECK-LABEL: p4hir.parser @p
18+
p4hir.parser @p(%arg0: !packet_in, %arg1: !p4hir.ref<!headers_t>)() {
19+
p4hir.state @start {
20+
p4hir.scope {
21+
%stack_ref = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
22+
%data_ref = p4hir.struct_field_ref %stack_ref["data"] : <!hs_2xH>
23+
%nextIndex_ref = p4hir.struct_field_ref %stack_ref["nextIndex"] : <!hs_2xH>
24+
%nextIndex = p4hir.read %nextIndex_ref : <!b32i>
25+
%elt_ref = p4hir.array_element_ref %data_ref[%nextIndex] : !p4hir.ref<!arr_2xH>, !b32i
26+
%hdr = p4hir.variable ["hdr"] : <!H>
27+
28+
// CHECK: p4hir.struct_field_ref {{.*}}["nextIndex"] : <!hs_2xH>
29+
// CHECK: p4hir.read {{.*}} : <!b32i>
30+
// CHECK: %[[size:.*]] = p4hir.const #int2_b32i
31+
// CHECK: %[[check:.*]] = p4hir.cmp(lt, {{.*}} : !b32i, %[[size]] : !b32i)
32+
// CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
33+
// CHECK: p4corelib.verify %[[check]] signalling %[[error]] : !error
34+
// CHECK: p4corelib.extract_header {{.*}} : <!H> from %arg0 : !p4corelib.packet_in
35+
p4hir.call_method @packet_in::@extract<[!H]> (%hdr) of %arg0 : !packet_in : (!p4hir.ref<!H>) -> ()
36+
37+
%val = p4hir.read %hdr : <!H>
38+
p4hir.assign %val, %elt_ref : <!H>
39+
40+
// CHECK: %[[one:.*]] = p4hir.const #int1_b32i
41+
// CHECK: p4hir.struct_field_ref {{.*}}["nextIndex"] : <!hs_2xH>
42+
// CHECK: %[[curr:.*]] = p4hir.read {{.*}} : <!b32i>
43+
// CHECK: %[[inc:.*]] = p4hir.binop(add, %[[curr]], %[[one]]) : !b32i
44+
// CHECK: p4hir.assign %[[inc]], {{.*}} : <!b32i>
45+
}
46+
p4hir.transition to @p::@accept
47+
}
48+
p4hir.state @accept {
49+
p4hir.parser_accept
50+
}
51+
p4hir.transition to @p::@start
52+
}
53+
}

‎test/Conversion/P4CoreLib/packet_in_hs.p4‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// RUN: p4mlir-translate --typeinference-only %s | sed 's/__corelib = \[\]/corelib/g' | p4mlir-opt --lower-to-p4corelib | FileCheck %s
2+
3+
@__corelib
4+
extern packet_in {
5+
void extract<T>(out T hdr);
6+
}
7+
8+
header H {
9+
bit<8> data;
10+
}
11+
12+
struct headers_t {
13+
H[2] stack;
14+
}
15+
16+
// CHECK-LABEL: p4hir.parser @p
17+
parser p(packet_in pkt, out headers_t hdr) {
18+
// CHECK: p4hir.state @start
19+
state start {
20+
// CHECK: p4hir.scope
21+
// CHECK: %[[stack_field_ref:.*]] = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
22+
// CHECK: %[[data_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["data"] : <!hs_2xH>
23+
// CHECK: %[[nextIndex_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
24+
// CHECK: %[[val:.*]] = p4hir.read %[[nextIndex_field_ref]] : <!b32i>
25+
// CHECK: %[[elt_ref:.*]] = p4hir.array_element_ref %[[data_field_ref]][%[[val]]] : !p4hir.ref<!arr_2xH>, !b32i
26+
// CHECK: %[[hdr_out_arg:.*]] = p4hir.variable ["hdr_out_arg"] : <!H>
27+
28+
// CHECK: %[[nextIndex_field_ref_0:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
29+
// CHECK: %[[val_1:.*]] = p4hir.read %[[nextIndex_field_ref_0]] : <!b32i>
30+
// CHECK: %[[c2_b32i:.*]] = p4hir.const #int2_b32i
31+
// CHECK: %[[lt:.*]] = p4hir.cmp(lt, %[[val_1]] : !b32i, %[[c2_b32i]] : !b32i)
32+
// CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
33+
// CHECK: p4corelib.verify %[[lt]] signalling %[[error]] : !error
34+
35+
// CHECK: p4corelib.extract_header %[[hdr_out_arg]] : <!H> from %arg0 : !p4corelib.packet_in
36+
// CHECK: %[[val_2:.*]] = p4hir.read %[[hdr_out_arg]] : <!H>
37+
// CHECK: p4hir.assign %[[val_2]], %[[elt_ref]] : <!H>
38+
39+
// CHECK: %[[c1_b32i:.*]] = p4hir.const #int1_b32i
40+
// CHECK: %[[nextIndex_field_ref_3:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
41+
// CHECK: %[[val_4:.*]] = p4hir.read %[[nextIndex_field_ref_3]] : <!b32i>
42+
// CHECK: %[[add:.*]] = p4hir.binop(add, %[[val_4]], %[[c1_b32i]]) : !b32i
43+
// CHECK: p4hir.assign %[[add]], %[[nextIndex_field_ref_3]] : <!b32i>
44+
pkt.extract(hdr.stack.next);
45+
transition accept;
46+
}
47+
}

‎test/Conversion/P4CoreLib/packet_in_hs_varbit.p4‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// RUN: p4mlir-translate --typeinference-only %s | sed 's/__corelib = \[\]/corelib/g' | p4mlir-opt --lower-to-p4corelib | FileCheck %s
2+
3+
@__corelib
4+
extern packet_in {
5+
void extract<T>(out T hdr, in bit<32> variableFieldSizeInBits);
6+
}
7+
8+
header H {
9+
bit<8> fixed_data;
10+
varbit<64> variable_data;
11+
}
12+
13+
struct headers_t {
14+
H[2] stack;
15+
}
16+
17+
// CHECK-LABEL: p4hir.parser @p
18+
parser p(packet_in pkt, out headers_t hdr) {
19+
// CHECK: p4hir.state @start
20+
state start {
21+
// CHECK: p4hir.scope
22+
// CHECK: %[[stack_field_ref:.*]] = p4hir.struct_field_ref %arg1["stack"] : <!headers_t>
23+
// CHECK: %[[data_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["data"] : <!hs_2xH>
24+
// CHECK: %[[nextIndex_field_ref:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
25+
// CHECK: %[[val:.*]] = p4hir.read %[[nextIndex_field_ref]] : <!b32i>
26+
// CHECK: %[[elt_ref:.*]] = p4hir.array_element_ref %[[data_field_ref]][%[[val]]] : !p4hir.ref<!arr_2xH>, !b32i
27+
// CHECK: %[[hdr_out_arg:.*]] = p4hir.variable ["hdr_out_arg"] : <!H>
28+
29+
// Bounds check inserted by lowering
30+
// CHECK: %[[nextIndex_field_ref_0:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
31+
// CHECK: %[[val_1:.*]] = p4hir.read %[[nextIndex_field_ref_0]] : <!b32i>
32+
// CHECK: %[[c2_b32i:.*]] = p4hir.const #int2_b32i
33+
// CHECK: %[[lt:.*]] = p4hir.cmp(lt, %[[val_1]] : !b32i, %[[c2_b32i]] : !b32i)
34+
// CHECK: %[[error:.*]] = p4hir.const #error_StackOutOfBounds
35+
// CHECK: p4corelib.verify %[[lt]] signalling %[[error]] : !error
36+
37+
// Variable-width extract lowered to p4corelib
38+
// CHECK: p4corelib.extract_header_variable %[[hdr_out_arg]]<{{.*}}> : <!H> from %arg0 : !p4corelib.packet_in
39+
40+
// nextIndex increment inserted by lowering
41+
// CHECK: %[[c1_b32i:.*]] = p4hir.const #int1_b32i
42+
// CHECK: %[[nextIndex_field_ref_3:.*]] = p4hir.struct_field_ref %[[stack_field_ref]]["nextIndex"] : <!hs_2xH>
43+
// CHECK: %[[val_4:.*]] = p4hir.read %[[nextIndex_field_ref_3]] : <!b32i>
44+
// CHECK: %[[add:.*]] = p4hir.binop(add, %[[val_4]], %[[c1_b32i]]) : !b32i
45+
// CHECK: p4hir.assign %[[add]], %[[nextIndex_field_ref_3]] : <!b32i>
46+
pkt.extract(hdr.stack.next, 32);
47+
transition accept;
48+
}
49+
}

0 commit comments

Comments
 (0)