sb->ste_build_tag_func = &dr_ste_v0_build_tnl_mpls_tag;
 }
 
-#define ICMP_TYPE_OFFSET_FIRST_DW              24
-#define ICMP_CODE_OFFSET_FIRST_DW              16
-#define ICMP_HEADER_DATA_OFFSET_SECOND_DW      0
+#define ICMP_TYPE_OFFSET_FIRST_DW      24
+#define ICMP_CODE_OFFSET_FIRST_DW      16
 
 static int
 dr_ste_v0_build_icmp_tag(struct mlx5dr_match_param *value,
                         u8 *tag)
 {
        struct mlx5dr_match_misc3 *misc_3 = &value->misc3;
-       u32 icmp_header_data;
+       u32 *icmp_header_data;
        int dw0_location;
        int dw1_location;
-       u32 icmp_type;
-       u32 icmp_code;
+       u8 *icmp_type;
+       u8 *icmp_code;
        bool is_ipv4;
 
        is_ipv4 = DR_MASK_IS_ICMPV4_SET(misc_3);
        if (is_ipv4) {
-               icmp_header_data        = misc_3->icmpv4_header_data;
-               icmp_type               = misc_3->icmpv4_type;
-               icmp_code               = misc_3->icmpv4_code;
+               icmp_header_data        = &misc_3->icmpv4_header_data;
+               icmp_type               = &misc_3->icmpv4_type;
+               icmp_code               = &misc_3->icmpv4_code;
                dw0_location            = sb->caps->flex_parser_id_icmp_dw0;
                dw1_location            = sb->caps->flex_parser_id_icmp_dw1;
        } else {
-               icmp_header_data        = misc_3->icmpv6_header_data;
-               icmp_type               = misc_3->icmpv6_type;
-               icmp_code               = misc_3->icmpv6_code;
+               icmp_header_data        = &misc_3->icmpv6_header_data;
+               icmp_type               = &misc_3->icmpv6_type;
+               icmp_code               = &misc_3->icmpv6_code;
                dw0_location            = sb->caps->flex_parser_id_icmpv6_dw0;
                dw1_location            = sb->caps->flex_parser_id_icmpv6_dw1;
        }
 
        switch (dw0_location) {
        case 4:
-               if (icmp_type) {
-                       MLX5_SET(ste_flex_parser_1, tag, flex_parser_4,
-                                (icmp_type << ICMP_TYPE_OFFSET_FIRST_DW));
-                       if (is_ipv4)
-                               misc_3->icmpv4_type = 0;
-                       else
-                               misc_3->icmpv6_type = 0;
-               }
+               MLX5_SET(ste_flex_parser_1, tag, flex_parser_4,
+                        (*icmp_type << ICMP_TYPE_OFFSET_FIRST_DW) |
+                        (*icmp_code << ICMP_TYPE_OFFSET_FIRST_DW));
 
-               if (icmp_code) {
-                       u32 cur_val = MLX5_GET(ste_flex_parser_1, tag,
-                                              flex_parser_4);
-                       MLX5_SET(ste_flex_parser_1, tag, flex_parser_4,
-                                cur_val | (icmp_code << ICMP_CODE_OFFSET_FIRST_DW));
-                       if (is_ipv4)
-                               misc_3->icmpv4_code = 0;
-                       else
-                               misc_3->icmpv6_code = 0;
-               }
+               *icmp_type = 0;
+               *icmp_code = 0;
                break;
        default:
                return -EINVAL;
 
        switch (dw1_location) {
        case 5:
-               if (icmp_header_data) {
-                       MLX5_SET(ste_flex_parser_1, tag, flex_parser_5,
-                                (icmp_header_data << ICMP_HEADER_DATA_OFFSET_SECOND_DW));
-                       if (is_ipv4)
-                               misc_3->icmpv4_header_data = 0;
-                       else
-                               misc_3->icmpv6_header_data = 0;
-               }
+               MLX5_SET(ste_flex_parser_1, tag, flex_parser_5,
+                        *icmp_header_data);
+               *icmp_header_data = 0;
                break;
        default:
                return -EINVAL;