static void mtk_iommu_tlb_flush_all(struct mtk_iommu_data *data)
 {
        for_each_m4u(data) {
+               if (pm_runtime_get_if_in_use(data->dev) <= 0)
+                       continue;
+
                writel_relaxed(F_INVLD_EN1 | F_INVLD_EN0,
                               data->base + data->plat_data->inv_sel_reg);
                writel_relaxed(F_ALL_INVLD, data->base + REG_MMU_INVALIDATE);
                wmb(); /* Make sure the tlb flush all done */
+
+               pm_runtime_put(data->dev);
        }
 }
 
                                           size_t granule,
                                           struct mtk_iommu_data *data)
 {
+       bool has_pm = !!data->dev->pm_domain;
        unsigned long flags;
        int ret;
        u32 tmp;
 
        for_each_m4u(data) {
+               if (has_pm) {
+                       if (pm_runtime_get_if_in_use(data->dev) <= 0)
+                               continue;
+               }
+
                spin_lock_irqsave(&data->tlb_lock, flags);
                writel_relaxed(F_INVLD_EN1 | F_INVLD_EN0,
                               data->base + data->plat_data->inv_sel_reg);
                /* Clear the CPE status */
                writel_relaxed(0, data->base + REG_MMU_CPE_DONE);
                spin_unlock_irqrestore(&data->tlb_lock, flags);
+
+               if (has_pm)
+                       pm_runtime_put(data->dev);
        }
 }
 
 {
        struct mtk_iommu_data *data = dev_iommu_priv_get(dev);
        struct mtk_iommu_domain *dom = to_mtk_domain(domain);
+       struct device *m4udev = data->dev;
        int ret;
 
        if (!data)
                return -ENODEV;
 
        if (!data->m4u_dom) { /* Initialize the M4U HW */
+               ret = pm_runtime_resume_and_get(m4udev);
+               if (ret < 0)
+                       return ret;
+
                ret = mtk_iommu_hw_init(data);
-               if (ret)
+               if (ret) {
+                       pm_runtime_put(m4udev);
                        return ret;
+               }
                data->m4u_dom = dom;
                writel(dom->cfg.arm_v7s_cfg.ttbr & MMU_PT_ADDR_MASK,
                       data->base + REG_MMU_PT_BASE_ADDR);
+
+               pm_runtime_put(m4udev);
        }
 
        mtk_iommu_config(data, dev, true);
        of_node_put(smicomm_node);
        data->smicomm_dev = &plarbdev->dev;
 
+       pm_runtime_enable(dev);
+
        link = device_link_add(data->smicomm_dev, dev,
                        DL_FLAG_STATELESS | DL_FLAG_PM_RUNTIME);
        if (!link) {
                dev_err(dev, "Unable link %s.\n", dev_name(data->smicomm_dev));
-               return -EINVAL;
+               goto out_runtime_disable;
        }
 
        platform_set_drvdata(pdev, data);
        iommu_device_sysfs_remove(&data->iommu);
 out_link_remove:
        device_link_remove(data->smicomm_dev, dev);
+out_runtime_disable:
+       pm_runtime_disable(dev);
        return ret;
 }
 
 
        clk_disable_unprepare(data->bclk);
        device_link_remove(data->smicomm_dev, &pdev->dev);
+       pm_runtime_disable(&pdev->dev);
        devm_free_irq(&pdev->dev, data->irq, data);
        component_master_del(&pdev->dev, &mtk_iommu_com_ops);
        return 0;
        void __iomem *base = data->base;
        int ret;
 
+       /* Avoid first resume to affect the default value of registers below. */
+       if (!m4u_dom)
+               return 0;
        ret = clk_prepare_enable(data->bclk);
        if (ret) {
                dev_err(data->dev, "Failed to enable clk(%d) in resume\n", ret);
        writel_relaxed(reg->int_main_control, base + REG_MMU_INT_MAIN_CONTROL);
        writel_relaxed(reg->ivrp_paddr, base + REG_MMU_IVRP_PADDR);
        writel_relaxed(reg->vld_pa_rng, base + REG_MMU_VLD_PA_RNG);
-       if (m4u_dom)
-               writel(m4u_dom->cfg.arm_v7s_cfg.ttbr & MMU_PT_ADDR_MASK,
-                      base + REG_MMU_PT_BASE_ADDR);
+       writel(m4u_dom->cfg.arm_v7s_cfg.ttbr & MMU_PT_ADDR_MASK, base + REG_MMU_PT_BASE_ADDR);
        return 0;
 }