#include <linux/of_address.h>
 #include <linux/of_irq.h>
 #include <linux/of_platform.h>
+#include <linux/pci.h>
 #include <linux/platform_device.h>
 #include <linux/pm_runtime.h>
 #include <linux/regmap.h>
 #define MTK_IOMMU_TYPE_MASK            (0x3 << 13)
 /* PM and clock always on. e.g. infra iommu */
 #define PM_CLK_AO                      BIT(15)
+#define IFA_IOMMU_PCIE_SUPPORT         BIT(16)
 
 #define MTK_IOMMU_HAS_FLAG_MASK(pdata, _x, mask)       \
                                ((((pdata)->flags) & (mask)) == (_x))
                                larb_mmu->mmu &= ~MTK_SMI_MMU_EN(portid);
                } else if (MTK_IOMMU_IS_TYPE(data->plat_data, MTK_IOMMU_TYPE_INFRA)) {
                        peri_mmuen_msk = BIT(portid);
-                       peri_mmuen = enable ? peri_mmuen_msk : 0;
+                       /* PCI dev has only one output id, enable the next writing bit for PCIe */
+                       if (dev_is_pci(dev))
+                               peri_mmuen_msk |= BIT(portid + 1);
 
+                       peri_mmuen = enable ? peri_mmuen_msk : 0;
                        ret = regmap_update_bits(data->pericfg, PERICFG_IOMMU_1,
                                                 peri_mmuen_msk, peri_mmuen);
                        if (ret)
                ret = component_master_add_with_match(dev, &mtk_iommu_com_ops, match);
                if (ret)
                        goto out_bus_set_null;
+       } else if (MTK_IOMMU_IS_TYPE(data->plat_data, MTK_IOMMU_TYPE_INFRA) &&
+                  MTK_IOMMU_HAS_FLAG(data->plat_data, IFA_IOMMU_PCIE_SUPPORT)) {
+#ifdef CONFIG_PCI
+               if (!iommu_present(&pci_bus_type)) {
+                       ret = bus_set_iommu(&pci_bus_type, &mtk_iommu_ops);
+                       if (ret) /* PCIe fail don't affect platform_bus. */
+                               goto out_list_del;
+               }
+#endif
        }
        return ret;
 
        if (MTK_IOMMU_IS_TYPE(data->plat_data, MTK_IOMMU_TYPE_MM)) {
                device_link_remove(data->smicomm_dev, &pdev->dev);
                component_master_del(&pdev->dev, &mtk_iommu_com_ops);
+       } else if (MTK_IOMMU_IS_TYPE(data->plat_data, MTK_IOMMU_TYPE_INFRA) &&
+                  MTK_IOMMU_HAS_FLAG(data->plat_data, IFA_IOMMU_PCIE_SUPPORT)) {
+#ifdef CONFIG_PCI
+               bus_set_iommu(&pci_bus_type, NULL);
+#endif
        }
        pm_runtime_disable(&pdev->dev);
        devm_free_irq(&pdev->dev, data->irq, data);