return byte_mask;
 }
 
+static u8 *mlx5dr_ste_get_tag(u8 *hw_ste_p)
+{
+       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
+
+       return hw_ste->tag;
+}
+
 void mlx5dr_ste_set_bit_mask(u8 *hw_ste_p, u8 *bit_mask)
 {
        struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
 
                mlx5dr_ste_set_bit_mask(ste_arr, sb->bit_mask);
 
-               ret = sb->ste_build_tag_func(value, sb, ste_arr);
+               ret = sb->ste_build_tag_func(value, sb, mlx5dr_ste_get_tag(ste_arr));
                if (ret)
                        return ret;
 
 
 static int dr_ste_build_eth_l2_src_des_tag(struct mlx5dr_match_param *value,
                                           struct mlx5dr_ste_build *sb,
-                                          u8 *hw_ste_p)
+                                          u8 *tag)
 {
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l2_src_dst, tag, dmac_47_16, spec, dmac_47_16);
        DR_STE_SET_TAG(eth_l2_src_dst, tag, dmac_15_0, spec, dmac_15_0);
 
 static int dr_ste_build_eth_l3_ipv6_dst_tag(struct mlx5dr_match_param *value,
                                            struct mlx5dr_ste_build *sb,
-                                           u8 *hw_ste_p)
+                                           u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l3_ipv6_dst, tag, dst_ip_127_96, spec, dst_ip_127_96);
        DR_STE_SET_TAG(eth_l3_ipv6_dst, tag, dst_ip_95_64, spec, dst_ip_95_64);
 
 static int dr_ste_build_eth_l3_ipv6_src_tag(struct mlx5dr_match_param *value,
                                            struct mlx5dr_ste_build *sb,
-                                           u8 *hw_ste_p)
+                                           u8 *tag)
 {
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l3_ipv6_src, tag, src_ip_127_96, spec, src_ip_127_96);
        DR_STE_SET_TAG(eth_l3_ipv6_src, tag, src_ip_95_64, spec, src_ip_95_64);
 
 static int dr_ste_build_eth_l3_ipv4_5_tuple_tag(struct mlx5dr_match_param *value,
                                                struct mlx5dr_ste_build *sb,
-                                               u8 *hw_ste_p)
+                                               u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l3_ipv4_5_tuple, tag, destination_address, spec, dst_ip_31_0);
        DR_STE_SET_TAG(eth_l3_ipv4_5_tuple, tag, source_address, spec, src_ip_31_0);
 }
 
 static int dr_ste_build_eth_l2_src_or_dst_tag(struct mlx5dr_match_param *value,
-                                             bool inner, u8 *hw_ste_p)
+                                             bool inner, u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = inner ? &value->inner : &value->outer;
        struct mlx5dr_match_misc *misc_spec = &value->misc;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l2_src, tag, first_vlan_id, spec, first_vid);
        DR_STE_SET_TAG(eth_l2_src, tag, first_cfi, spec, first_cfi);
 
 static int dr_ste_build_eth_l2_src_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l2_src, tag, smac_47_16, spec, smac_47_16);
        DR_STE_SET_TAG(eth_l2_src, tag, smac_15_0, spec, smac_15_0);
 
-       return dr_ste_build_eth_l2_src_or_dst_tag(value, sb->inner, hw_ste_p);
+       return dr_ste_build_eth_l2_src_or_dst_tag(value, sb->inner, tag);
 }
 
 void mlx5dr_ste_build_eth_l2_src(struct mlx5dr_ste_build *sb,
 
 static int dr_ste_build_eth_l2_dst_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l2_dst, tag, dmac_47_16, spec, dmac_47_16);
        DR_STE_SET_TAG(eth_l2_dst, tag, dmac_15_0, spec, dmac_15_0);
 
-       return dr_ste_build_eth_l2_src_or_dst_tag(value, sb->inner, hw_ste_p);
+       return dr_ste_build_eth_l2_src_or_dst_tag(value, sb->inner, tag);
 }
 
 void mlx5dr_ste_build_eth_l2_dst(struct mlx5dr_ste_build *sb,
 
 static int dr_ste_build_eth_l2_tnl_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc *misc = &value->misc;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l2_tnl, tag, dmac_47_16, spec, dmac_47_16);
        DR_STE_SET_TAG(eth_l2_tnl, tag, dmac_15_0, spec, dmac_15_0);
 
 static int dr_ste_build_eth_l3_ipv4_misc_tag(struct mlx5dr_match_param *value,
                                             struct mlx5dr_ste_build *sb,
-                                            u8 *hw_ste_p)
+                                            u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l3_ipv4_misc, tag, time_to_live, spec, ttl_hoplimit);
 
 
 static int dr_ste_build_ipv6_l3_l4_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
        struct mlx5dr_match_spec *spec = sb->inner ? &value->inner : &value->outer;
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(eth_l4, tag, dst_port, spec, tcp_dport);
        DR_STE_SET_TAG(eth_l4, tag, src_port, spec, tcp_sport);
 
 static int dr_ste_build_empty_always_hit_tag(struct mlx5dr_match_param *value,
                                             struct mlx5dr_ste_build *sb,
-                                            u8 *hw_ste_p)
+                                            u8 *tag)
 {
        return 0;
 }
 
 static int dr_ste_build_mpls_tag(struct mlx5dr_match_param *value,
                                 struct mlx5dr_ste_build *sb,
-                                u8 *hw_ste_p)
+                                u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc2 *misc2_mask = &value->misc2;
-       u8 *tag = hw_ste->tag;
 
        if (sb->inner)
                DR_STE_SET_MPLS_TAG(mpls, misc2_mask, inner, tag);
 
 static int dr_ste_build_gre_tag(struct mlx5dr_match_param *value,
                                struct mlx5dr_ste_build *sb,
-                               u8 *hw_ste_p)
+                               u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct  mlx5dr_match_misc *misc = &value->misc;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(gre, tag, gre_protocol, misc, gre_protocol);
 
 
 static int dr_ste_build_flex_parser_0_tag(struct mlx5dr_match_param *value,
                                          struct mlx5dr_ste_build *sb,
-                                         u8 *hw_ste_p)
+                                         u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc2 *misc_2_mask = &value->misc2;
-       u8 *tag = hw_ste->tag;
 
        if (DR_STE_IS_OUTER_MPLS_OVER_GRE_SET(misc_2_mask)) {
                DR_STE_SET_TAG(flex_parser_0, tag, parser_3_label,
 
 static int dr_ste_build_flex_parser_1_tag(struct mlx5dr_match_param *value,
                                          struct mlx5dr_ste_build *sb,
-                                         u8 *hw_ste_p)
+                                         u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc3 *misc_3 = &value->misc3;
-       u8 *tag = hw_ste->tag;
        u32 icmp_header_data;
        int dw0_location;
        int dw1_location;
 
 static int dr_ste_build_general_purpose_tag(struct mlx5dr_match_param *value,
                                            struct mlx5dr_ste_build *sb,
-                                           u8 *hw_ste_p)
+                                           u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc2 *misc_2_mask = &value->misc2;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(general_purpose, tag, general_purpose_lookup_field,
                       misc_2_mask, metadata_reg_a);
 
 static int dr_ste_build_eth_l4_misc_tag(struct mlx5dr_match_param *value,
                                        struct mlx5dr_ste_build *sb,
-                                       u8 *hw_ste_p)
+                                       u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc3 *misc3 = &value->misc3;
-       u8 *tag = hw_ste->tag;
 
        if (sb->inner) {
                DR_STE_SET_TAG(eth_l4_misc, tag, seq_num, misc3, inner_tcp_seq_num);
 static int
 dr_ste_build_flex_parser_tnl_vxlan_gpe_tag(struct mlx5dr_match_param *value,
                                           struct mlx5dr_ste_build *sb,
-                                          u8 *hw_ste_p)
+                                          u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc3 *misc3 = &value->misc3;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(flex_parser_tnl_vxlan_gpe, tag,
                       outer_vxlan_gpe_flags, misc3,
 static int
 dr_ste_build_flex_parser_tnl_geneve_tag(struct mlx5dr_match_param *value,
                                        struct mlx5dr_ste_build *sb,
-                                       u8 *hw_ste_p)
+                                       u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc *misc = &value->misc;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(flex_parser_tnl_geneve, tag,
                       geneve_protocol_type, misc, geneve_protocol_type);
 
 static int dr_ste_build_register_0_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc2 *misc2 = &value->misc2;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(register_0, tag, register_0_h, misc2, metadata_reg_c_0);
        DR_STE_SET_TAG(register_0, tag, register_0_l, misc2, metadata_reg_c_1);
 
 static int dr_ste_build_register_1_tag(struct mlx5dr_match_param *value,
                                       struct mlx5dr_ste_build *sb,
-                                      u8 *hw_ste_p)
+                                      u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc2 *misc2 = &value->misc2;
-       u8 *tag = hw_ste->tag;
 
        DR_STE_SET_TAG(register_1, tag, register_2_h, misc2, metadata_reg_c_4);
        DR_STE_SET_TAG(register_1, tag, register_2_l, misc2, metadata_reg_c_5);
 
 static int dr_ste_build_src_gvmi_qpn_tag(struct mlx5dr_match_param *value,
                                         struct mlx5dr_ste_build *sb,
-                                        u8 *hw_ste_p)
+                                        u8 *tag)
 {
-       struct dr_hw_ste_format *hw_ste = (struct dr_hw_ste_format *)hw_ste_p;
        struct mlx5dr_match_misc *misc = &value->misc;
        struct mlx5dr_cmd_vport_cap *vport_cap;
        struct mlx5dr_domain *dmn = sb->dmn;
        struct mlx5dr_cmd_caps *caps;
        u8 *bit_mask = sb->bit_mask;
-       u8 *tag = hw_ste->tag;
        bool source_gvmi_set;
 
        DR_STE_SET_TAG(src_gvmi_qp, tag, source_qp, misc, source_sqn);