return pe;
 }
 
-long pnv_npu_set_window(struct pnv_ioda_pe *npe, int num,
+static long pnv_npu_unset_window(struct iommu_table_group *table_group,
+               int num);
+
+static long pnv_npu_set_window(struct iommu_table_group *table_group, int num,
                struct iommu_table *tbl)
 {
+       struct pnv_ioda_pe *npe = container_of(table_group, struct pnv_ioda_pe,
+                       table_group);
        struct pnv_phb *phb = npe->phb;
        int64_t rc;
        const unsigned long size = tbl->it_indirect_levels ?
 
        /* NPU has just one TVE so if there is another table, remove it first */
        if (npe->table_group.tables[num2])
-               pnv_npu_unset_window(npe, num2);
+               pnv_npu_unset_window(&npe->table_group, num2);
 
        pe_info(npe, "Setting up window %llx..%llx pg=%lx\n",
                        start_addr, start_addr + win_size - 1,
        return 0;
 }
 
-long pnv_npu_unset_window(struct pnv_ioda_pe *npe, int num)
+static long pnv_npu_unset_window(struct iommu_table_group *table_group, int num)
 {
+       struct pnv_ioda_pe *npe = container_of(table_group, struct pnv_ioda_pe,
+                       table_group);
        struct pnv_phb *phb = npe->phb;
        int64_t rc;
 
        if (!gpe)
                return;
 
-       rc = pnv_npu_set_window(npe, 0, gpe->table_group.tables[0]);
+       rc = pnv_npu_set_window(&npe->table_group, 0,
+                       gpe->table_group.tables[0]);
 
        /*
         * NVLink devices use the same TCE table configuration as
        if (phb->type != PNV_PHB_NPU_NVLINK || !npe->pdev)
                return -EINVAL;
 
-       rc = pnv_npu_unset_window(npe, 0);
+       rc = pnv_npu_unset_window(&npe->table_group, 0);
        if (rc != OPAL_SUCCESS)
                return rc;
 
        }
 }
 
+#ifdef CONFIG_IOMMU_API
 /* Switch ownership from platform code to external user (e.g. VFIO) */
-void pnv_npu_take_ownership(struct pnv_ioda_pe *npe)
+static void pnv_npu_take_ownership(struct iommu_table_group *table_group)
 {
+       struct pnv_ioda_pe *npe = container_of(table_group, struct pnv_ioda_pe,
+                       table_group);
        struct pnv_phb *phb = npe->phb;
        int64_t rc;
 
         * if it was enabled at the moment of ownership change.
         */
        if (npe->table_group.tables[0]) {
-               pnv_npu_unset_window(npe, 0);
+               pnv_npu_unset_window(&npe->table_group, 0);
                return;
        }
 
        pnv_pci_ioda2_tce_invalidate_entire(npe->phb, false);
 }
 
+static struct iommu_table_group_ops pnv_pci_npu_ops = {
+       .set_window = pnv_npu_set_window,
+       .unset_window = pnv_npu_unset_window,
+       .take_ownership = pnv_npu_take_ownership,
+};
+
 struct pnv_ioda_pe *pnv_pci_npu_setup_iommu(struct pnv_ioda_pe *npe)
 {
        struct pnv_phb *phb = npe->phb;
        if (!gpe || !gpdev)
                return NULL;
 
+       npe->table_group.ops = &pnv_pci_npu_ops;
+
        list_for_each_entry(npdev, &pbus->devices, bus_list) {
                gptmp = pnv_pci_get_gpu_dev(npdev);
 
 
        return gpe;
 }
+#endif /* !CONFIG_IOMMU_API */
 
 /*
  * NPU2 ATS
 
 static long pnv_pci_ioda2_npu_set_window(struct iommu_table_group *table_group,
                int num, struct iommu_table *tbl)
 {
+       struct pnv_ioda_pe *npe = gpe_table_group_to_npe(table_group);
        long ret = pnv_pci_ioda2_set_window(table_group, num, tbl);
 
        if (ret)
                return ret;
 
-       ret = pnv_npu_set_window(gpe_table_group_to_npe(table_group), num, tbl);
+       ret = npe->table_group.ops->set_window(&npe->table_group, num, tbl);
        if (ret)
                pnv_pci_ioda2_unset_window(table_group, num);
 
                struct iommu_table_group *table_group,
                int num)
 {
+       struct pnv_ioda_pe *npe = gpe_table_group_to_npe(table_group);
        long ret = pnv_pci_ioda2_unset_window(table_group, num);
 
        if (ret)
                return ret;
 
-       return pnv_npu_unset_window(gpe_table_group_to_npe(table_group), num);
+       return npe->table_group.ops->unset_window(&npe->table_group, num);
 }
 
 static void pnv_ioda2_npu_take_ownership(struct iommu_table_group *table_group)
 {
-       pnv_npu_take_ownership(gpe_table_group_to_npe(table_group));
+       struct pnv_ioda_pe *npe = gpe_table_group_to_npe(table_group);
+
+       npe->table_group.ops->take_ownership(&npe->table_group);
        pnv_ioda2_take_ownership(table_group);
 }
 
 
 extern void pnv_npu_try_dma_set_bypass(struct pci_dev *gpdev, bool bypass);
 extern void pnv_pci_ioda2_tce_invalidate_entire(struct pnv_phb *phb, bool rm);
 extern struct pnv_ioda_pe *pnv_pci_npu_setup_iommu(struct pnv_ioda_pe *npe);
-extern long pnv_npu_set_window(struct pnv_ioda_pe *npe, int num,
-               struct iommu_table *tbl);
-extern long pnv_npu_unset_window(struct pnv_ioda_pe *npe, int num);
-extern void pnv_npu_take_ownership(struct pnv_ioda_pe *npe);
-extern void pnv_npu_release_ownership(struct pnv_ioda_pe *npe);
 
 /* pci-ioda-tce.c */
 #define POWERNV_IOMMU_DEFAULT_LEVELS   1