return mode;
 }
 
+static bool ath9k_htc_setpower(struct ath9k_htc_priv *priv,
+                              enum ath9k_power_mode mode)
+{
+       bool ret;
+
+       mutex_lock(&priv->htc_pm_lock);
+       ret = ath9k_hw_setpower(priv->ah, mode);
+       mutex_unlock(&priv->htc_pm_lock);
+
+       return ret;
+}
+
+void ath9k_htc_ps_wakeup(struct ath9k_htc_priv *priv)
+{
+       mutex_lock(&priv->htc_pm_lock);
+       if (++priv->ps_usecount != 1)
+               goto unlock;
+       ath9k_hw_setpower(priv->ah, ATH9K_PM_AWAKE);
+
+unlock:
+       mutex_unlock(&priv->htc_pm_lock);
+}
+
+void ath9k_htc_ps_restore(struct ath9k_htc_priv *priv)
+{
+       mutex_lock(&priv->htc_pm_lock);
+       if (--priv->ps_usecount != 0)
+               goto unlock;
+
+       if (priv->ps_enabled)
+               ath9k_hw_setpower(priv->ah, ATH9K_PM_NETWORK_SLEEP);
+unlock:
+       mutex_unlock(&priv->htc_pm_lock);
+}
+
+void ath9k_ps_work(struct work_struct *work)
+{
+       struct ath9k_htc_priv *priv =
+               container_of(work, struct ath9k_htc_priv,
+                            ps_work);
+       ath9k_htc_setpower(priv, ATH9K_PM_AWAKE);
+
+       /* The chip wakes up after receiving the first beacon
+          while network sleep is enabled. For the driver to
+          be in sync with the hw, set the chip to awake and
+          only then set it to sleep.
+        */
+       ath9k_htc_setpower(priv, ATH9K_PM_NETWORK_SLEEP);
+}
+
 static int ath9k_htc_set_channel(struct ath9k_htc_priv *priv,
                                 struct ieee80211_hw *hw,
                                 struct ath9k_channel *hchan)
 
        /* Fiddle around with fastcc later on, for now just use full reset */
        fastcc = false;
-
+       ath9k_htc_ps_wakeup(priv);
        htc_stop(priv->htc);
        WMI_CMD(WMI_DISABLE_INTR_CMDID);
        WMI_CMD(WMI_DRAIN_TXQ_ALL_CMDID);
                ath_print(common, ATH_DBG_FATAL,
                          "Unable to reset channel (%u Mhz) "
                          "reset status %d\n", channel->center_freq, ret);
+               ath9k_htc_ps_restore(priv);
                goto err;
        }
 
 
        priv->op_flags &= ~OP_FULL_RESET;
 err:
+       ath9k_htc_ps_restore(priv);
        return ret;
 }
 
 
        short_cal_interval = ATH_STA_SHORT_CALINTERVAL;
 
+       /* Only calibrate if awake */
+       if (ah->power_mode != ATH9K_PM_AWAKE)
+               goto set_timer;
+
        /* Long calibration runs independently of short calibration. */
        if ((timestamp - common->ani.longcal_timer) >= ATH_LONG_CALINTERVAL) {
                longcal = true;
 
        /* Skip all processing if there's nothing to do. */
        if (longcal || shortcal || aniflag) {
+
+               ath9k_htc_ps_wakeup(priv);
+
                /* Call ANI routine if necessary */
                if (aniflag)
                        ath9k_hw_ani_monitor(ah, ah->curchan);
                                  ah->curchan->channelFlags,
                                  common->ani.noise_floor);
                }
+
+               ath9k_htc_ps_restore(priv);
        }
 
+set_timer:
        /*
        * Set timer interval based on previous results.
        * The interval must be the shortest necessary to satisfy ANI,
                return;
        }
 
+       ath9k_htc_ps_wakeup(priv);
        htc_stop(priv->htc);
        WMI_CMD(WMI_DISABLE_INTR_CMDID);
        WMI_CMD(WMI_DRAIN_TXQ_ALL_CMDID);
        ath9k_hw_phy_disable(ah);
        ath9k_hw_disable(ah);
        ath9k_hw_configpcipowersave(ah, 1, 1);
-       ath9k_hw_setpower(ah, ATH9K_PM_FULL_SLEEP);
+       ath9k_htc_ps_restore(priv);
+       ath9k_htc_setpower(priv, ATH9K_PM_FULL_SLEEP);
 
+       cancel_work_sync(&priv->ps_work);
        cancel_delayed_work_sync(&priv->ath9k_ani_work);
        cancel_delayed_work_sync(&priv->ath9k_aggr_work);
        cancel_delayed_work_sync(&priv->ath9k_led_blink_work);
                goto out;
        }
 
+       ath9k_htc_ps_wakeup(priv);
        memset(&hvif, 0, sizeof(struct ath9k_htc_target_vif));
        memcpy(&hvif.myaddr, vif->addr, ETH_ALEN);
 
 
        priv->vif = vif;
 out:
+       ath9k_htc_ps_restore(priv);
        mutex_unlock(&priv->mutex);
        return ret;
 }
                }
 
        }
+       if (changed & IEEE80211_CONF_CHANGE_PS) {
+               if (conf->flags & IEEE80211_CONF_PS) {
+                       ath9k_htc_setpower(priv, ATH9K_PM_NETWORK_SLEEP);
+                       priv->ps_enabled = true;
+               } else {
+                       priv->ps_enabled = false;
+                       cancel_work_sync(&priv->ps_work);
+                       ath9k_htc_setpower(priv, ATH9K_PM_AWAKE);
+               }
+       }
 
        if (changed & IEEE80211_CONF_CHANGE_MONITOR) {
                if (conf->flags & IEEE80211_CONF_MONITOR) {
 
        mutex_lock(&priv->mutex);
 
+       ath9k_htc_ps_wakeup(priv);
        changed_flags &= SUPPORTED_FILTERS;
        *total_flags &= SUPPORTED_FILTERS;
 
        ath_print(ath9k_hw_common(priv->ah), ATH_DBG_CONFIG,
                  "Set HW RX filter: 0x%x\n", rfilt);
 
+       ath9k_htc_ps_restore(priv);
        mutex_unlock(&priv->mutex);
 }
 
 
        mutex_lock(&priv->mutex);
        ath_print(common, ATH_DBG_CONFIG, "Set HW Key\n");
+       ath9k_htc_ps_wakeup(priv);
 
        switch (cmd) {
        case SET_KEY:
                ret = -EINVAL;
        }
 
+       ath9k_htc_ps_restore(priv);
        mutex_unlock(&priv->mutex);
 
        return ret;
        struct ath_common *common = ath9k_hw_common(ah);
 
        mutex_lock(&priv->mutex);
+       ath9k_htc_ps_wakeup(priv);
 
        if (changed & BSS_CHANGED_ASSOC) {
                common->curaid = bss_conf->assoc ?
                        ath_start_ani(priv);
                } else {
                        priv->op_flags &= ~OP_ASSOCIATED;
+                       cancel_work_sync(&priv->ps_work);
                        cancel_delayed_work_sync(&priv->ath9k_ani_work);
                }
        }
                ath9k_hw_init_global_settings(ah);
        }
 
+       ath9k_htc_ps_restore(priv);
        mutex_unlock(&priv->mutex);
 }
 
 {
        struct ath9k_htc_priv *priv = hw->priv;
 
+       ath9k_htc_ps_wakeup(priv);
        mutex_lock(&priv->mutex);
        ath9k_hw_reset_tsf(priv->ah);
        mutex_unlock(&priv->mutex);
+       ath9k_htc_ps_restore(priv);
 }
 
 static int ath9k_htc_ampdu_action(struct ieee80211_hw *hw,
        spin_lock_bh(&priv->beacon_lock);
        priv->op_flags |= OP_SCANNING;
        spin_unlock_bh(&priv->beacon_lock);
+       cancel_work_sync(&priv->ps_work);
        cancel_delayed_work_sync(&priv->ath9k_ani_work);
        mutex_unlock(&priv->mutex);
 }
 {
        struct ath9k_htc_priv *priv = hw->priv;
 
+       ath9k_htc_ps_wakeup(priv);
        mutex_lock(&priv->mutex);
        spin_lock_bh(&priv->beacon_lock);
        priv->op_flags &= ~OP_SCANNING;
        priv->op_flags |= OP_FULL_RESET;
        ath_start_ani(priv);
        mutex_unlock(&priv->mutex);
+       ath9k_htc_ps_restore(priv);
 }
 
 static int ath9k_htc_set_rts_threshold(struct ieee80211_hw *hw, u32 value)