Add comparison operator, handle Operation_Relation

Signed-off-by: Anton Korobeynikov <[email protected]>
p4lang · Jan 31, 2025 · 84d735f · 84d735f
1 parent 2968abe
commit 84d735f
Show file tree

Hide file tree

Showing 4 changed files with 135 additions and 0 deletions.
diff --git a/include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td b/include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td
@@ -285,4 +285,43 @@ def BinOp : P4HIR_Op<"binop", [Pure,
   let hasVerifier = 0;
 }
 
+def CmpOpKind_LT : I32EnumAttrCase<"Lt", 1, "lt">;
+def CmpOpKind_LE : I32EnumAttrCase<"Le", 2, "le">;
+def CmpOpKind_GT : I32EnumAttrCase<"Gt", 3, "gt">;
+def CmpOpKind_GE : I32EnumAttrCase<"Ge", 4, "ge">;
+def CmpOpKind_EQ : I32EnumAttrCase<"Eq", 5, "eq">;
+def CmpOpKind_NE : I32EnumAttrCase<"Ne", 6, "ne">;
+
+def CmpOpKind : I32EnumAttr<
+    "CmpOpKind",
+    "compare operation kind",
+    [CmpOpKind_LT, CmpOpKind_LE, CmpOpKind_GT,
+     CmpOpKind_GE, CmpOpKind_EQ, CmpOpKind_NE]> {
+  let cppNamespace = "::P4::P4MLIR::P4HIR";
+}
+
+def CmpOp : P4HIR_Op<"cmp", [Pure, SameTypeOperands]> {
+  let summary = "Compare values two values and produce a boolean result";
+  let description = [{
+    `p4hir.cmp` compares two input operands of the same type and produces a
+    `p4hir.bool` result. The kinds of comparison available are:
+    [lt,gt,ge,eq,ne]
+
+    ```mlir
+    %7 = p4hir.cmp(gt, %1, %2) : !p4hir.bit<32>, !p4hir.bool
+    ```
+  }];
+
+  let results = (outs BooleanType:$result);
+  let arguments = (ins Arg<CmpOpKind, "cmp kind">:$kind,
+                       AnyP4Type:$lhs, AnyP4Type:$rhs);
+
+  let assemblyFormat = [{
+    `(` $kind `,` $lhs `,` $rhs  `)` `:` type($lhs) `,` type($result) attr-dict
+  }];
+
+  // Already covered by the traits
+  let hasVerifier = 0;
+}
+
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_OPS_TD
diff --git a/test/Dialect/P4HIR/cmp.mlir b/test/Dialect/P4HIR/cmp.mlir
@@ -0,0 +1,15 @@
+// RUN: p4mlir-opt %s | FileCheck %s
+
+// No need to check stuff. If it parses, it's fine.
+// CHECK: module
+module {
+  %0 = p4hir.const #p4hir.int<42> : !p4hir.bit<8>
+  %1 = p4hir.const #p4hir.int<42> : !p4hir.bit<8>
+
+  %2 = p4hir.cmp(eq, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+  %3 = p4hir.cmp(ne, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+  %4 = p4hir.cmp(lt, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+  %5 = p4hir.cmp(le, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+  %6 = p4hir.cmp(ge, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+  %7 = p4hir.cmp(gt, %0, %1) : !p4hir.bit<8>, !p4hir.bool
+}
diff --git a/test/Translate/Ops/cmp.p4 b/test/Translate/Ops/cmp.p4
@@ -0,0 +1,52 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+
+// CHECK-LABEL:   module
+// CHECK-NEXT:    %[[VAL_0:.*]] = p4hir.alloca !p4hir.bool ["res"] : !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_1:.*]] = p4hir.const #p4hir.int<1> : !p4hir.bit<10>
+// CHECK:         %[[VAL_2:.*]] = p4hir.cast(%[[VAL_1]] : !p4hir.bit<10>) : !p4hir.bit<10>
+// CHECK:         %[[VAL_3:.*]] = p4hir.alloca !p4hir.bit<10> ["lhs", init] : !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         p4hir.store %[[VAL_2]], %[[VAL_3]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_4:.*]] = p4hir.const #p4hir.int<2> : !p4hir.bit<10>
+// CHECK:         %[[VAL_5:.*]] = p4hir.cast(%[[VAL_4]] : !p4hir.bit<10>) : !p4hir.bit<10>
+// CHECK:         %[[VAL_6:.*]] = p4hir.alloca !p4hir.bit<10> ["rhs", init] : !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         p4hir.store %[[VAL_5]], %[[VAL_6]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_7:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_8:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_9:.*]] = p4hir.cmp(eq, %[[VAL_7]], %[[VAL_8]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_9]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_10:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_11:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_12:.*]] = p4hir.cmp(ne, %[[VAL_10]], %[[VAL_11]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_12]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_13:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_14:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_15:.*]] = p4hir.cmp(lt, %[[VAL_13]], %[[VAL_14]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_15]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_16:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_17:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_18:.*]] = p4hir.cmp(gt, %[[VAL_16]], %[[VAL_17]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_18]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_19:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_20:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_21:.*]] = p4hir.cmp(le, %[[VAL_19]], %[[VAL_20]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_21]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+// CHECK:         %[[VAL_22:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_23:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_24:.*]] = p4hir.cmp(ge, %[[VAL_22]], %[[VAL_23]]) : !p4hir.bit<10>, !p4hir.bool
+// CHECK:         p4hir.store %[[VAL_24]], %[[VAL_0]] : !p4hir.bool, !p4hir.ref<!p4hir.bool>
+
+action cmp() {
+    bool res;
+
+    bit<10> lhs = 1;
+    bit<10> rhs = 2;
+
+    res = lhs == rhs;
+    res = lhs != rhs;
+    res = lhs < rhs;
+    res = lhs > rhs;
+    res = lhs <= rhs;
+    res = lhs >= rhs;
+}
diff --git a/tools/p4mlir-translate/translate.cpp b/tools/p4mlir-translate/translate.cpp
@@ -292,6 +292,14 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
     HANDLE_IN_POSTORDER(BAnd)
     HANDLE_IN_POSTORDER(BXor)
 
+    // Comparisons
+    HANDLE_IN_POSTORDER(Equ)
+    HANDLE_IN_POSTORDER(Neq)
+    HANDLE_IN_POSTORDER(Leq)
+    HANDLE_IN_POSTORDER(Lss)
+    HANDLE_IN_POSTORDER(Grt)
+    HANDLE_IN_POSTORDER(Geq)
+
     HANDLE_IN_POSTORDER(Cast)
     HANDLE_IN_POSTORDER(Declaration_Variable)
 
@@ -302,6 +310,7 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
 
     mlir::Value emitUnOp(const P4::IR::Operation_Unary *unop, P4HIR::UnaryOpKind kind);
     mlir::Value emitBinOp(const P4::IR::Operation_Binary *binop, P4HIR::BinOpKind kind);
+    mlir::Value emitCmp(const P4::IR::Operation_Relation *relop, P4HIR::CmpOpKind kind);
 };
 
 bool P4TypeConverter::preorder(const P4::IR::Type_Bits *type) {
@@ -480,6 +489,11 @@ mlir::Value P4HIRConverter::emitBinOp(const P4::IR::Operation_Binary *binop,
     return builder.create<P4HIR::BinOp>(getLoc(builder, binop), kind, getValue(binop->left),
                                         getValue(binop->right));
 }
+mlir::Value P4HIRConverter::emitCmp(const P4::IR::Operation_Relation *relop,
+                                    P4HIR::CmpOpKind kind) {
+    return builder.create<P4HIR::CmpOp>(getLoc(builder, relop), kind, getValue(relop->left),
+                                        getValue(relop->right));
+}
 
 #define CONVERT_UNOP(Node, Kind)                                  \
     void P4HIRConverter::postorder(const P4::IR::Node *node) {    \
@@ -513,6 +527,21 @@ CONVERT_BINOP(BXor, Xor);
 
 #undef CONVERT_BINOP
 
+#define CONVERT_CMP(Node, Kind)                                \
+    void P4HIRConverter::postorder(const P4::IR::Node *node) { \
+        ConversionTracer trace("Converting ", node);           \
+        setValue(node, emitCmp(node, P4HIR::CmpOpKind::Kind)); \
+    }
+
+CONVERT_CMP(Equ, Eq);
+CONVERT_CMP(Neq, Ne);
+CONVERT_CMP(Lss, Lt);
+CONVERT_CMP(Leq, Le);
+CONVERT_CMP(Grt, Gt);
+CONVERT_CMP(Geq, Ge);
+
+#undef CONVERT_CMP
+
 bool P4HIRConverter::preorder(const P4::IR::AssignmentStatement *assign) {
     ConversionTracer trace("Converting ", assign);