phy_write_mmd(phydev, MDIO_MMD_PCS, offsets[i],
                                      mac[(i * 2) + 1] | (mac[(i * 2)] << 8));
 
-               /* Enable WOL function for 1588 */
-               if (phydev->drv->phy_id == ATH8031_PHY_ID) {
-                       ret = phy_modify_mmd(phydev, MDIO_MMD_PCS,
-                                            AT803X_PHY_MMD3_WOL_CTRL,
-                                            0, AT803X_WOL_EN);
-                       if (ret)
-                               return ret;
-               }
                /* Enable WOL interrupt */
                ret = phy_modify(phydev, AT803X_INTR_ENABLE, 0, AT803X_INTR_ENABLE_WOL);
                if (ret)
                        return ret;
        } else {
-               /* Disable WoL function for 1588 */
-               if (phydev->drv->phy_id == ATH8031_PHY_ID) {
-                       ret = phy_modify_mmd(phydev, MDIO_MMD_PCS,
-                                            AT803X_PHY_MMD3_WOL_CTRL,
-                                            AT803X_WOL_EN, 0);
-                       if (ret)
-                               return ret;
-               }
                /* Disable WOL interrupt */
                ret = phy_modify(phydev, AT803X_INTR_ENABLE, AT803X_INTR_ENABLE_WOL, 0);
                if (ret)
        return at803x_config_init(phydev);
 }
 
+static int at8031_set_wol(struct phy_device *phydev,
+                         struct ethtool_wolinfo *wol)
+{
+       int ret;
+
+       /* First setup MAC address and enable WOL interrupt */
+       ret = at803x_set_wol(phydev, wol);
+       if (ret)
+               return ret;
+
+       if (wol->wolopts & WAKE_MAGIC)
+               /* Enable WOL function for 1588 */
+               ret = phy_modify_mmd(phydev, MDIO_MMD_PCS,
+                                    AT803X_PHY_MMD3_WOL_CTRL,
+                                    0, AT803X_WOL_EN);
+       else
+               /* Disable WoL function for 1588 */
+               ret = phy_modify_mmd(phydev, MDIO_MMD_PCS,
+                                    AT803X_PHY_MMD3_WOL_CTRL,
+                                    AT803X_WOL_EN, 0);
+
+       return ret;
+}
+
 static int qca83xx_config_init(struct phy_device *phydev)
 {
        u8 switch_revision;
        .config_init            = at8031_config_init,
        .config_aneg            = at803x_config_aneg,
        .soft_reset             = genphy_soft_reset,
-       .set_wol                = at803x_set_wol,
+       .set_wol                = at8031_set_wol,
        .get_wol                = at803x_get_wol,
        .suspend                = at803x_suspend,
        .resume                 = at803x_resume,