};
 #endif
 
+static void iwl_mvm_reset_phy_ctxts(struct iwl_mvm *mvm)
+{
+       int i;
+
+       memset(mvm->phy_ctxts, 0, sizeof(mvm->phy_ctxts));
+       for (i = 0; i < NUM_PHY_CTX; i++) {
+               mvm->phy_ctxts[i].id = i;
+               mvm->phy_ctxts[i].ref = 0;
+       }
+}
+
 int iwl_mvm_mac_setup_register(struct iwl_mvm *mvm)
 {
        struct ieee80211_hw *hw = mvm->hw;
 
        hw->sta_data_size = sizeof(struct iwl_mvm_sta);
        hw->vif_data_size = sizeof(struct iwl_mvm_vif);
-       hw->chanctx_data_size = sizeof(struct iwl_mvm_phy_ctxt);
+       hw->chanctx_data_size = sizeof(u16);
 
        hw->wiphy->interface_modes = BIT(NL80211_IFTYPE_STATION) |
                BIT(NL80211_IFTYPE_P2P_CLIENT) |
                hw->wiphy->n_addresses++;
        }
 
+       iwl_mvm_reset_phy_ctxts(mvm);
+
        /* we create the 802.11 header and a max-length SSID element */
        hw->wiphy->max_scan_ie_len =
                mvm->fw->ucode_capa.max_probe_length - 24 - 34;
        iwl_mvm_te_clear_data(mvm, &mvmvif->time_event_data);
        spin_unlock_bh(&mvm->time_event_lock);
 
-       if (vif->type != NL80211_IFTYPE_P2P_DEVICE)
-               mvmvif->phy_ctxt = NULL;
+       mvmvif->phy_ctxt = NULL;
 }
 
 static void iwl_mvm_restart_cleanup(struct iwl_mvm *mvm)
                mvm->hw, IEEE80211_IFACE_ITER_RESUME_ALL,
                iwl_mvm_cleanup_iterator, mvm);
 
+       mvm->p2p_device_vif = NULL;
+
+       iwl_mvm_reset_phy_ctxts(mvm);
        memset(mvm->fw_key_table, 0, sizeof(mvm->fw_key_table));
        memset(mvm->sta_drained, 0, sizeof(mvm->sta_drained));
 
        iwl_mvm_power_update_mode(mvm, vif);
 }
 
+static struct iwl_mvm_phy_ctxt *iwl_mvm_get_free_phy_ctxt(struct iwl_mvm *mvm)
+{
+       u16 i;
+
+       lockdep_assert_held(&mvm->mutex);
+
+       for (i = 0; i < NUM_PHY_CTX; i++)
+               if (!mvm->phy_ctxts[i].ref)
+                       return &mvm->phy_ctxts[i];
+
+       IWL_ERR(mvm, "No available PHY context\n");
+       return NULL;
+}
+
 static int iwl_mvm_mac_add_interface(struct ieee80211_hw *hw,
                                     struct ieee80211_vif *vif)
 {
                struct ieee80211_channel *chan;
                struct cfg80211_chan_def chandef;
 
-               mvmvif->phy_ctxt = &mvm->phy_ctxt_roc;
+               mvmvif->phy_ctxt = iwl_mvm_get_free_phy_ctxt(mvm);
+               if (!mvmvif->phy_ctxt) {
+                       ret = -ENOSPC;
+                       goto out_remove_mac;
+               }
 
                /*
                 * The channel used here isn't relevant as it's
  out_unbind:
        iwl_mvm_binding_remove_vif(mvm, vif);
  out_remove_phy:
-       iwl_mvm_phy_ctxt_remove(mvm, mvmvif->phy_ctxt);
+       iwl_mvm_phy_ctxt_unref(mvm, mvmvif->phy_ctxt);
  out_remove_mac:
        mvmvif->phy_ctxt = NULL;
        iwl_mvm_mac_ctxt_remove(mvm, vif);
                mvm->p2p_device_vif = NULL;
                iwl_mvm_rm_bcast_sta(mvm, &mvmvif->bcast_sta);
                iwl_mvm_binding_remove_vif(mvm, vif);
-               iwl_mvm_phy_ctxt_remove(mvm, mvmvif->phy_ctxt);
+               iwl_mvm_phy_ctxt_unref(mvm, mvmvif->phy_ctxt);
                mvmvif->phy_ctxt = NULL;
        }
 
                       enum ieee80211_roc_type type)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
+       struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
        struct cfg80211_chan_def chandef;
        int ret;
 
        mutex_lock(&mvm->mutex);
 
        cfg80211_chandef_create(&chandef, channel, NL80211_CHAN_NO_HT);
-       ret = iwl_mvm_phy_ctxt_changed(mvm, &mvm->phy_ctxt_roc,
+       ret = iwl_mvm_phy_ctxt_changed(mvm, mvmvif->phy_ctxt,
                                       &chandef, 1, 1);
 
        /* Schedule the time events */
                               struct ieee80211_chanctx_conf *ctx)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
-       struct iwl_mvm_phy_ctxt *phy_ctxt = (void *)ctx->drv_priv;
+       u16 *phy_ctxt_id = (u16 *)ctx->drv_priv;
+       struct iwl_mvm_phy_ctxt *phy_ctxt;
        int ret;
 
+       IWL_DEBUG_MAC80211(mvm, "Add PHY context\n");
+
        mutex_lock(&mvm->mutex);
+       phy_ctxt = iwl_mvm_get_free_phy_ctxt(mvm);
+       if (!phy_ctxt) {
+               ret = -ENOSPC;
+               goto out;
+       }
 
-       IWL_DEBUG_MAC80211(mvm, "Add PHY context\n");
        ret = iwl_mvm_phy_ctxt_add(mvm, phy_ctxt, &ctx->def,
                                   ctx->rx_chains_static,
                                   ctx->rx_chains_dynamic);
+       if (ret) {
+               IWL_ERR(mvm, "Failed to add PHY context\n");
+               goto out;
+       }
+
+       *phy_ctxt_id = phy_ctxt->id;
+out:
        mutex_unlock(&mvm->mutex);
        return ret;
 }
                                   struct ieee80211_chanctx_conf *ctx)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
-       struct iwl_mvm_phy_ctxt *phy_ctxt = (void *)ctx->drv_priv;
+       u16 *phy_ctxt_id = (u16 *)ctx->drv_priv;
+       struct iwl_mvm_phy_ctxt *phy_ctxt = &mvm->phy_ctxts[*phy_ctxt_id];
 
        mutex_lock(&mvm->mutex);
-       iwl_mvm_phy_ctxt_remove(mvm, phy_ctxt);
+       iwl_mvm_phy_ctxt_unref(mvm, phy_ctxt);
        mutex_unlock(&mvm->mutex);
 }
 
                                   u32 changed)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
-       struct iwl_mvm_phy_ctxt *phy_ctxt = (void *)ctx->drv_priv;
+       u16 *phy_ctxt_id = (u16 *)ctx->drv_priv;
+       struct iwl_mvm_phy_ctxt *phy_ctxt = &mvm->phy_ctxts[*phy_ctxt_id];
 
        mutex_lock(&mvm->mutex);
        iwl_mvm_phy_ctxt_changed(mvm, phy_ctxt, &ctx->def,
                                      struct ieee80211_chanctx_conf *ctx)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
-       struct iwl_mvm_phy_ctxt *phyctx = (void *)ctx->drv_priv;
+       u16 *phy_ctxt_id = (u16 *)ctx->drv_priv;
+       struct iwl_mvm_phy_ctxt *phy_ctxt = &mvm->phy_ctxts[*phy_ctxt_id];
        struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
        int ret;
 
        mutex_lock(&mvm->mutex);
 
-       mvmvif->phy_ctxt = phyctx;
+       mvmvif->phy_ctxt = phy_ctxt;
 
        switch (vif->type) {
        case NL80211_IFTYPE_AP:
 
        return ret;
 }
 
-
-struct phy_ctx_used_data {
-       unsigned long used[BITS_TO_LONGS(NUM_PHY_CTX)];
-};
-
-static void iwl_mvm_phy_ctx_used_iter(struct ieee80211_hw *hw,
-                                     struct ieee80211_chanctx_conf *ctx,
-                                     void *_data)
-{
-       struct phy_ctx_used_data *data = _data;
-       struct iwl_mvm_phy_ctxt *phy_ctxt = (void *)ctx->drv_priv;
-
-       __set_bit(phy_ctxt->id, data->used);
-}
-
 /*
  * Send a command to add a PHY context based on the current HW configuration.
  */
                         struct cfg80211_chan_def *chandef,
                         u8 chains_static, u8 chains_dynamic)
 {
-       struct phy_ctx_used_data data = {
-               .used = { },
-       };
-
-       /*
-        * If this is a regular PHY context (not the ROC one)
-        * skip the ROC PHY context's ID.
-        */
-       if (ctxt != &mvm->phy_ctxt_roc)
-               __set_bit(mvm->phy_ctxt_roc.id, data.used);
+       int ret;
 
+       WARN_ON(ctxt->ref);
        lockdep_assert_held(&mvm->mutex);
-       ctxt->color++;
 
-       if (!test_bit(IWL_MVM_STATUS_IN_HW_RESTART, &mvm->status)) {
-               ieee80211_iter_chan_contexts_atomic(
-                       mvm->hw, iwl_mvm_phy_ctx_used_iter, &data);
+       ctxt->channel = chandef->chan;
+       ret = iwl_mvm_phy_ctxt_apply(mvm, ctxt, chandef,
+                                    chains_static, chains_dynamic,
+                                    FW_CTXT_ACTION_ADD, 0);
 
-               ctxt->id = find_first_zero_bit(data.used, NUM_PHY_CTX);
-               if (WARN_ONCE(ctxt->id == NUM_PHY_CTX,
-                             "Failed to init PHY context - no free ID!\n"))
-                       return -EIO;
-       }
+       if (!ret)
+               ctxt->ref = 1;
+       return ret;
+}
 
-       ctxt->channel = chandef->chan;
-       return iwl_mvm_phy_ctxt_apply(mvm, ctxt, chandef,
-                                     chains_static, chains_dynamic,
-                                     FW_CTXT_ACTION_ADD, 0);
+/*
+ * Update the number of references to the given PHY context. This is valid only
+ * in case the PHY context was already created, i.e., its reference count > 0.
+ */
+void iwl_mvm_phy_ctxt_ref(struct iwl_mvm *mvm, struct iwl_mvm_phy_ctxt *ctxt)
+{
+       WARN_ON(!ctxt->ref);
+       lockdep_assert_held(&mvm->mutex);
+
+       ctxt->ref++;
 }
 
 /*
  * Once the command is sent, regardless of success or failure, the context is
  * marked as invalid
  */
-void iwl_mvm_phy_ctxt_remove(struct iwl_mvm *mvm, struct iwl_mvm_phy_ctxt *ctxt)
+static void iwl_mvm_phy_ctxt_remove(struct iwl_mvm *mvm,
+                                   struct iwl_mvm_phy_ctxt *ctxt)
 {
        struct iwl_phy_context_cmd cmd;
        int ret;
        ret = iwl_mvm_send_cmd_pdu(mvm, PHY_CONTEXT_CMD, CMD_SYNC,
                                   sizeof(struct iwl_phy_context_cmd),
                                   &cmd);
+       ctxt->channel = NULL;
        if (ret)
                IWL_ERR(mvm, "Failed to send PHY remove: ctxt id=%d\n",
                        ctxt->id);
 }
+
+void iwl_mvm_phy_ctxt_unref(struct iwl_mvm *mvm, struct iwl_mvm_phy_ctxt *ctxt)
+{
+       lockdep_assert_held(&mvm->mutex);
+
+       ctxt->ref--;
+       if (ctxt->ref != 0)
+               return;
+
+       return iwl_mvm_phy_ctxt_remove(mvm, ctxt);
+}