return mh->modify_hdr;
 }
 
+char *
+mlx5e_mod_hdr_alloc(struct mlx5_core_dev *mdev, int namespace,
+                   struct mlx5e_tc_mod_hdr_acts *mod_hdr_acts)
+{
+       int new_num_actions, max_hw_actions;
+       size_t new_sz, old_sz;
+       void *ret;
+
+       if (mod_hdr_acts->num_actions < mod_hdr_acts->max_actions)
+               goto out;
+
+       max_hw_actions = mlx5e_mod_hdr_max_actions(mdev, namespace);
+       new_num_actions = min(max_hw_actions,
+                             mod_hdr_acts->actions ?
+                             mod_hdr_acts->max_actions * 2 : 1);
+       if (mod_hdr_acts->max_actions == new_num_actions)
+               return ERR_PTR(-ENOSPC);
+
+       new_sz = MLX5_MH_ACT_SZ * new_num_actions;
+       old_sz = mod_hdr_acts->max_actions * MLX5_MH_ACT_SZ;
+
+       ret = krealloc(mod_hdr_acts->actions, new_sz, GFP_KERNEL);
+       if (!ret)
+               return ERR_PTR(-ENOMEM);
+
+       memset(ret + old_sz, 0, new_sz - old_sz);
+       mod_hdr_acts->actions = ret;
+       mod_hdr_acts->max_actions = new_num_actions;
+
+out:
+       return mod_hdr_acts->actions + (mod_hdr_acts->num_actions * MLX5_MH_ACT_SZ);
+}
+
+void
+mlx5e_mod_hdr_dealloc(struct mlx5e_tc_mod_hdr_acts *mod_hdr_acts)
+{
+       kfree(mod_hdr_acts->actions);
+       mod_hdr_acts->actions = NULL;
+       mod_hdr_acts->num_actions = 0;
+       mod_hdr_acts->max_actions = 0;
+}
+
+char *
+mlx5e_mod_hdr_get_item(struct mlx5e_tc_mod_hdr_acts *mod_hdr_acts, int pos)
+{
+       return mod_hdr_acts->actions + (pos * MLX5_MH_ACT_SZ);
+}
 
        struct flow_action *flow_action = &flow_rule->action;
        struct mlx5_core_dev *mdev = ct_priv->dev;
        struct flow_action_entry *act;
-       size_t action_size;
        char *modact;
        int err, i;
 
-       action_size = MLX5_UN_SZ_BYTES(set_add_copy_action_in_auto);
-
        flow_action_for_each(i, act, flow_action) {
                switch (act->id) {
                case FLOW_ACTION_MANGLE: {
-                       err = alloc_mod_hdr_actions(mdev, ct_priv->ns_type,
-                                                   mod_acts);
-                       if (err)
-                               return err;
-
-                       modact = mod_acts->actions +
-                                mod_acts->num_actions * action_size;
+                       modact = mlx5e_mod_hdr_alloc(mdev, ct_priv->ns_type, mod_acts);
+                       if (IS_ERR(modact))
+                               return PTR_ERR(modact);
 
                        err = mlx5_tc_ct_parse_mangle_to_mod_act(act, modact);
                        if (err)
                attr->modify_hdr = mlx5e_mod_hdr_get(*mh);
        }
 
-       dealloc_mod_hdr_actions(&mod_acts);
+       mlx5e_mod_hdr_dealloc(&mod_acts);
        return 0;
 
 err_mapping:
-       dealloc_mod_hdr_actions(&mod_acts);
+       mlx5e_mod_hdr_dealloc(&mod_acts);
        mlx5_put_label_mapping(ct_priv, attr->ct_attr.ct_labels_id);
        return err;
 }
        }
        pre_ct->miss_rule = rule;
 
-       dealloc_mod_hdr_actions(&pre_mod_acts);
+       mlx5e_mod_hdr_dealloc(&pre_mod_acts);
        kvfree(spec);
        return 0;
 
 err_flow_rule:
        mlx5_modify_header_dealloc(dev, pre_ct->modify_hdr);
 err_mapping:
-       dealloc_mod_hdr_actions(&pre_mod_acts);
+       mlx5e_mod_hdr_dealloc(&pre_mod_acts);
        kvfree(spec);
        return err;
 }
        }
 
        attr->ct_attr.ct_flow = ct_flow;
-       dealloc_mod_hdr_actions(&pre_mod_acts);
+       mlx5e_mod_hdr_dealloc(&pre_mod_acts);
 
        return ct_flow->pre_ct_rule;
 
 err_insert_orig:
        mlx5_modify_header_dealloc(priv->mdev, pre_ct_attr->modify_hdr);
 err_mapping:
-       dealloc_mod_hdr_actions(&pre_mod_acts);
+       mlx5e_mod_hdr_dealloc(&pre_mod_acts);
        mlx5_chains_put_chain_mapping(ct_priv->chains, ct_flow->chain_mapping);
 err_get_chain:
        kfree(ct_flow->pre_ct_attr);
 
 #include "lag/mp.h"
 
 #define nic_chains(priv) ((priv)->fs.tc.chains)
-#define MLX5_MH_ACT_SZ MLX5_UN_SZ_BYTES(set_add_copy_action_in_auto)
 
 #define MLX5E_TC_TABLE_NUM_GROUPS 4
 #define MLX5E_TC_TABLE_MAX_GROUP_SIZE BIT(18)
        char *modact;
        int err;
 
-       err = alloc_mod_hdr_actions(mdev, ns, mod_hdr_acts);
-       if (err)
-               return err;
-
-       modact = mod_hdr_acts->actions +
-                (mod_hdr_acts->num_actions * MLX5_MH_ACT_SZ);
+       modact = mlx5e_mod_hdr_alloc(mdev, ns, mod_hdr_acts);
+       if (IS_ERR(modact))
+               return PTR_ERR(modact);
 
        /* Firmware has 5bit length field and 0 means 32bits */
        if (mlen == 32)
        int mlen = mlx5e_tc_attr_to_reg_mappings[type].mlen;
        char *modact;
 
-       modact = mod_hdr_acts->actions + (act_id * MLX5_MH_ACT_SZ);
+       modact = mlx5e_mod_hdr_get_item(mod_hdr_acts, act_id);
 
        /* Firmware has 5bit length field and 0 means 32bits */
        if (mlen == 32)
 
        if (attr->action & MLX5_FLOW_CONTEXT_ACTION_MOD_HDR) {
                err = mlx5e_attach_mod_hdr(priv, flow, parse_attr);
-               dealloc_mod_hdr_actions(&parse_attr->mod_hdr_acts);
+               mlx5e_mod_hdr_dealloc(&parse_attr->mod_hdr_acts);
                if (err)
                        return err;
        }
        mlx5_tc_ct_match_del(get_ct_priv(priv), &flow->attr->ct_attr);
 
        if (attr->action & MLX5_FLOW_CONTEXT_ACTION_MOD_HDR) {
-               dealloc_mod_hdr_actions(&attr->parse_attr->mod_hdr_acts);
+               mlx5e_mod_hdr_dealloc(&attr->parse_attr->mod_hdr_acts);
                if (vf_tun && attr->modify_hdr)
                        mlx5_modify_header_dealloc(priv->mdev, attr->modify_hdr);
                else
                                struct netlink_ext_ack *extack)
 {
        struct pedit_headers *set_masks, *add_masks, *set_vals, *add_vals;
-       int i, action_size, first, last, next_z;
        void *headers_c, *headers_v, *action, *vals_p;
        u32 *s_masks_p, *a_masks_p, s_mask, a_mask;
        struct mlx5e_tc_mod_hdr_acts *mod_acts;
-       struct mlx5_fields *f;
        unsigned long mask, field_mask;
-       int err;
+       int i, first, last, next_z;
+       struct mlx5_fields *f;
        u8 cmd;
 
        mod_acts = &parse_attr->mod_hdr_acts;
        set_vals = &hdrs[0].vals;
        add_vals = &hdrs[1].vals;
 
-       action_size = MLX5_UN_SZ_BYTES(set_add_copy_action_in_auto);
-
        for (i = 0; i < ARRAY_SIZE(fields); i++) {
                bool skip;
 
                        return -EOPNOTSUPP;
                }
 
-               err = alloc_mod_hdr_actions(priv->mdev, namespace, mod_acts);
-               if (err) {
+               action = mlx5e_mod_hdr_alloc(priv->mdev, namespace, mod_acts);
+               if (IS_ERR(action)) {
                        NL_SET_ERR_MSG_MOD(extack,
                                           "too many pedit actions, can't offload");
                        mlx5_core_warn(priv->mdev,
                                       "mlx5: parsed %d pedit actions, can't do more\n",
                                       mod_acts->num_actions);
-                       return err;
+                       return PTR_ERR(action);
                }
 
-               action = mod_acts->actions +
-                        (mod_acts->num_actions * action_size);
                MLX5_SET(set_action_in, action, action_type, cmd);
                MLX5_SET(set_action_in, action, field, f->field);
 
        return 0;
 }
 
-static int mlx5e_flow_namespace_max_modify_action(struct mlx5_core_dev *mdev,
-                                                 int namespace)
-{
-       if (namespace == MLX5_FLOW_NAMESPACE_FDB) /* FDB offloading */
-               return MLX5_CAP_ESW_FLOWTABLE_FDB(mdev, max_modify_header_actions);
-       else /* namespace is MLX5_FLOW_NAMESPACE_KERNEL - NIC offloading */
-               return MLX5_CAP_FLOWTABLE_NIC_RX(mdev, max_modify_header_actions);
-}
-
-int alloc_mod_hdr_actions(struct mlx5_core_dev *mdev,
-                         int namespace,
-                         struct mlx5e_tc_mod_hdr_acts *mod_hdr_acts)
-{
-       int action_size, new_num_actions, max_hw_actions;
-       size_t new_sz, old_sz;
-       void *ret;
-
-       if (mod_hdr_acts->num_actions < mod_hdr_acts->max_actions)
-               return 0;
-
-       action_size = MLX5_UN_SZ_BYTES(set_add_copy_action_in_auto);
-
-       max_hw_actions = mlx5e_flow_namespace_max_modify_action(mdev,
-                                                               namespace);
-       new_num_actions = min(max_hw_actions,
-                             mod_hdr_acts->actions ?
-                             mod_hdr_acts->max_actions * 2 : 1);
-       if (mod_hdr_acts->max_actions == new_num_actions)
-               return -ENOSPC;
-
-       new_sz = action_size * new_num_actions;
-       old_sz = mod_hdr_acts->max_actions * action_size;
-       ret = krealloc(mod_hdr_acts->actions, new_sz, GFP_KERNEL);
-       if (!ret)
-               return -ENOMEM;
-
-       memset(ret + old_sz, 0, new_sz - old_sz);
-       mod_hdr_acts->actions = ret;
-       mod_hdr_acts->max_actions = new_num_actions;
-
-       return 0;
-}
-
-void dealloc_mod_hdr_actions(struct mlx5e_tc_mod_hdr_acts *mod_hdr_acts)
-{
-       kfree(mod_hdr_acts->actions);
-       mod_hdr_acts->actions = NULL;
-       mod_hdr_acts->num_actions = 0;
-       mod_hdr_acts->max_actions = 0;
-}
-
 static const struct pedit_headers zero_masks = {};
 
 static int
                goto out_err;
        }
 
-       if (!mlx5e_flow_namespace_max_modify_action(priv->mdev, namespace)) {
+       if (!mlx5e_mod_hdr_max_actions(priv->mdev, namespace)) {
                NL_SET_ERR_MSG_MOD(extack,
                                   "The pedit offload action is not supported");
                goto out_err;
        return 0;
 
 out_dealloc_parsed_actions:
-       dealloc_mod_hdr_actions(&parse_attr->mod_hdr_acts);
+       mlx5e_mod_hdr_dealloc(&parse_attr->mod_hdr_acts);
        return err;
 }
 
                return 0;
 
        attr->action &= ~MLX5_FLOW_CONTEXT_ACTION_MOD_HDR;
-       dealloc_mod_hdr_actions(&parse_attr->mod_hdr_acts);
+       mlx5e_mod_hdr_dealloc(&parse_attr->mod_hdr_acts);
 
        if (ns_type != MLX5_FLOW_NAMESPACE_FDB)
                return 0;
 
 err_free:
        flow_flag_set(flow, FAILED);
-       dealloc_mod_hdr_actions(&parse_attr->mod_hdr_acts);
+       mlx5e_mod_hdr_dealloc(&parse_attr->mod_hdr_acts);
        mlx5e_flow_put(priv, flow);
 out:
        return err;