#include <linux/mutex.h>
 #include <linux/of.h>
 #include <linux/platform_device.h>
+#include <linux/pm_domain.h>
 #include <linux/pm_runtime.h>
 #include <linux/spinlock_types.h>
 #include <linux/spinlock.h>
        if (!res->line_num)
                return -EINVAL;
 
-       if (res->has_pd)
-               vfe->genpd = camss->genpd[id];
+       if (res->has_pd) {
+               vfe->genpd = dev_pm_domain_attach_by_id(camss->dev, id);
+               if (IS_ERR(vfe->genpd))
+                       return PTR_ERR(vfe->genpd);
+       }
 
        vfe->line_num = res->line_num;
        vfe->ops->subdev_init(dev, vfe);
        return 0;
 }
 
+/*
+ * msm_vfe_genpd_cleanup - Cleanup VFE genpd linkages
+ * @vfe: VFE device
+ */
+void msm_vfe_genpd_cleanup(struct vfe_device *vfe)
+{
+       if (vfe->genpd_link)
+               device_link_del(vfe->genpd_link);
+
+       if (vfe->genpd)
+               dev_pm_domain_detach(vfe->genpd, true);
+}
+
 /*
  * vfe_link_setup - Setup VFE connections
  * @entity: Pointer to media entity structure
 
 
 static int camss_configure_pd(struct camss *camss)
 {
+       const struct camss_resources *res = camss->res;
        struct device *dev = camss->dev;
+       int vfepd_num;
        int i;
        int ret;
 
        if (camss->genpd_num == 1)
                return 0;
 
-       camss->genpd = devm_kmalloc_array(dev, camss->genpd_num,
-                                         sizeof(*camss->genpd), GFP_KERNEL);
-       if (!camss->genpd)
-               return -ENOMEM;
+       /* count the # of VFEs which have flagged power-domain */
+       for (vfepd_num = i = 0; i < camss->vfe_total_num; i++) {
+               if (res->vfe_res[i].has_pd)
+                       vfepd_num++;
+       }
 
-       camss->genpd_link = devm_kmalloc_array(dev, camss->genpd_num,
-                                              sizeof(*camss->genpd_link),
-                                              GFP_KERNEL);
-       if (!camss->genpd_link)
-               return -ENOMEM;
+       /*
+        * If the number of power-domains is greater than the number of VFEs
+        * then the additional power-domain is for the entire CAMSS block.
+        */
+       if (!(camss->genpd_num > vfepd_num))
+               return 0;
 
        /*
         * VFE power domains are in the beginning of the list, and while all
         * power domains should be attached, only if TITAN_TOP power domain is
         * found in the list, it should be linked over here.
         */
-       for (i = 0; i < camss->genpd_num; i++) {
-               camss->genpd[i] = dev_pm_domain_attach_by_id(camss->dev, i);
-               if (IS_ERR(camss->genpd[i])) {
-                       ret = PTR_ERR(camss->genpd[i]);
-                       goto fail_pm;
-               }
+       camss->genpd = dev_pm_domain_attach_by_id(camss->dev, camss->genpd_num - 1);
+       if (IS_ERR(camss->genpd)) {
+               ret = PTR_ERR(camss->genpd);
+               goto fail_pm;
        }
-
-       if (i > camss->res->vfe_num) {
-               camss->genpd_link[i - 1] = device_link_add(camss->dev, camss->genpd[i - 1],
-                                                          DL_FLAG_STATELESS | DL_FLAG_PM_RUNTIME |
-                                                          DL_FLAG_RPM_ACTIVE);
-               if (!camss->genpd_link[i - 1]) {
-                       ret = -EINVAL;
-                       goto fail_pm;
-               }
+       camss->genpd_link = device_link_add(camss->dev, camss->genpd,
+                                           DL_FLAG_STATELESS | DL_FLAG_PM_RUNTIME |
+                                           DL_FLAG_RPM_ACTIVE);
+       if (!camss->genpd_link) {
+               ret = -EINVAL;
+               goto fail_pm;
        }
 
        return 0;
 
 fail_pm:
-       for (--i ; i >= 0; i--)
-               dev_pm_domain_detach(camss->genpd[i], true);
+       dev_pm_domain_detach(camss->genpd, true);
 
        return ret;
 }
        return 0;
 }
 
-static void camss_genpd_cleanup(struct camss *camss)
+static void camss_genpd_subdevice_cleanup(struct camss *camss)
 {
        int i;
 
+       for (i = 0; i < camss->vfe_total_num; i++)
+               msm_vfe_genpd_cleanup(&camss->vfe[i]);
+}
+
+static void camss_genpd_cleanup(struct camss *camss)
+{
        if (camss->genpd_num == 1)
                return;
 
-       if (camss->genpd_num > camss->res->vfe_num)
-               device_link_del(camss->genpd_link[camss->genpd_num - 1]);
+       camss_genpd_subdevice_cleanup(camss);
+
+       if (camss->genpd_link)
+               device_link_del(camss->genpd_link);
 
-       for (i = 0; i < camss->genpd_num; i++)
-               dev_pm_domain_detach(camss->genpd[i], true);
+       dev_pm_domain_detach(camss->genpd, true);
 }
 
 /*