}
 
 /*
- * For the NPU we want to point the TCE table at the same table as the
- * real PCI device.
+ * Enables 32 bit DMA on NPU.
  */
-static void pnv_npu_disable_bypass(struct pnv_ioda_pe *npe)
+static void pnv_npu_dma_set_32(struct pnv_ioda_pe *npe)
 {
        struct pnv_phb *phb = npe->phb;
        struct pci_dev *gpdev;
 }
 
 /*
- * Enable/disable bypass mode on the NPU. The NPU only supports one
+ * Enables bypass mode on the NPU. The NPU only supports one
  * window per link, so bypass needs to be explicitly enabled or
  * disabled. Unlike for a PHB3 bypass and non-bypass modes can't be
  * active at the same time.
  */
-int pnv_npu_dma_set_bypass(struct pnv_ioda_pe *npe, bool enable)
+static int pnv_npu_dma_set_bypass(struct pnv_ioda_pe *npe)
 {
        struct pnv_phb *phb = npe->phb;
        int64_t rc = 0;
+       phys_addr_t top = memblock_end_of_DRAM();
 
        if (phb->type != PNV_PHB_NPU || !npe->pdev)
                return -EINVAL;
 
-       if (enable) {
-               /* Enable the bypass window */
-               phys_addr_t top = memblock_end_of_DRAM();
-
-               npe->tce_bypass_base = 0;
-               top = roundup_pow_of_two(top);
-               dev_info(&npe->pdev->dev, "Enabling bypass for PE %d\n",
-                        npe->pe_number);
-               rc = opal_pci_map_pe_dma_window_real(phb->opal_id,
-                                       npe->pe_number, npe->pe_number,
-                                       npe->tce_bypass_base, top);
-       } else {
-               /*
-                * Disable the bypass window by replacing it with the
-                * TCE32 window.
-                */
-               pnv_npu_disable_bypass(npe);
-       }
+       /* Enable the bypass window */
+
+       top = roundup_pow_of_two(top);
+       dev_info(&npe->pdev->dev, "Enabling bypass for PE %d\n",
+                       npe->pe_number);
+       rc = opal_pci_map_pe_dma_window_real(phb->opal_id,
+                       npe->pe_number, npe->pe_number,
+                       0 /* bypass base */, top);
 
        return rc;
 }
 
-int pnv_npu_dma_set_mask(struct pci_dev *npdev, u64 dma_mask)
+void pnv_npu_try_dma_set_bypass(struct pci_dev *gpdev, bool bypass)
 {
-       struct pci_controller *hose = pci_bus_to_host(npdev->bus);
-       struct pnv_phb *phb = hose->private_data;
-       struct pci_dn *pdn = pci_get_pdn(npdev);
-       struct pnv_ioda_pe *npe, *gpe;
-       struct pci_dev *gpdev;
-       uint64_t top;
-       bool bypass = false;
+       int i;
+       struct pnv_phb *phb;
+       struct pci_dn *pdn;
+       struct pnv_ioda_pe *npe;
+       struct pci_dev *npdev;
 
-       if (WARN_ON(!pdn || pdn->pe_number == IODA_INVALID_PE))
-               return -ENXIO;
+       for (i = 0; ; ++i) {
+               npdev = pnv_pci_get_npu_dev(gpdev, i);
 
-       /* We only do bypass if it's enabled on the linked device */
-       npe = &phb->ioda.pe_array[pdn->pe_number];
-       gpe = get_gpu_pci_dev_and_pe(npe, &gpdev);
-       if (!gpe)
-               return -ENODEV;
+               if (!npdev)
+                       break;
 
-       if (gpe->tce_bypass_enabled) {
-               top = gpe->tce_bypass_base + memblock_end_of_DRAM() - 1;
-               bypass = (dma_mask >= top);
-       }
+               pdn = pci_get_pdn(npdev);
+               if (WARN_ON(!pdn || pdn->pe_number == IODA_INVALID_PE))
+                       return;
 
-       if (bypass)
-               dev_info(&npdev->dev, "Using 64-bit DMA iommu bypass\n");
-       else
-               dev_info(&npdev->dev, "Using 32-bit DMA via iommu\n");
+               phb = pci_bus_to_host(npdev->bus)->private_data;
 
-       pnv_npu_dma_set_bypass(npe, bypass);
-       *npdev->dev.dma_mask = dma_mask;
+               /* We only do bypass if it's enabled on the linked device */
+               npe = &phb->ioda.pe_array[pdn->pe_number];
 
-       return 0;
+               if (bypass) {
+                       dev_info(&npdev->dev,
+                                       "Using 64-bit DMA iommu bypass\n");
+                       pnv_npu_dma_set_bypass(npe);
+               } else {
+                       dev_info(&npdev->dev, "Using 32-bit DMA via iommu\n");
+                       pnv_npu_dma_set_32(npe);
+               }
+       }
 }
 
        struct pnv_ioda_pe *pe;
        uint64_t top;
        bool bypass = false;
-       struct pci_dev *linked_npu_dev;
-       int i;
 
        if (WARN_ON(!pdn || pdn->pe_number == IODA_INVALID_PE))
                return -ENODEV;;
        *pdev->dev.dma_mask = dma_mask;
 
        /* Update peer npu devices */
-       if (pe->flags & PNV_IODA_PE_PEER)
-               for (i = 0; i < PNV_IODA_MAX_PEER_PES; i++) {
-                       if (!pe->peers[i])
-                               continue;
-
-                       linked_npu_dev = pe->peers[i]->pdev;
-                       if (dma_get_mask(&linked_npu_dev->dev) != dma_mask)
-                               dma_set_mask(&linked_npu_dev->dev, dma_mask);
-               }
+       pnv_npu_try_dma_set_bypass(pdev, bypass);
 
        return 0;
 }
                        enable_bypass = dma_get_mask(&pe->pdev->dev) ==
                                DMA_BIT_MASK(64);
                        pnv_npu_init_dma_pe(pe);
-                       pnv_npu_dma_set_bypass(pe, enable_bypass);
                }
        }
 }
        .shutdown               = pnv_pci_ioda_shutdown,
 };
 
+static int pnv_npu_dma_set_mask(struct pci_dev *npdev, u64 dma_mask)
+{
+       dev_err_once(&npdev->dev,
+                       "%s operation unsupported for NVLink devices\n",
+                       __func__);
+       return -EPERM;
+}
+
 static const struct pci_controller_ops pnv_npu_ioda_controller_ops = {
        .dma_dev_setup          = pnv_pci_dma_dev_setup,
 #ifdef CONFIG_PCI_MSI
        phb->freeze_pe = pnv_ioda_freeze_pe;
        phb->unfreeze_pe = pnv_ioda_unfreeze_pe;
 
-       /* Setup TCEs */
-       phb->dma_dev_setup = pnv_pci_ioda_dma_dev_setup;
-
        /* Setup MSI support */
        pnv_pci_init_ioda_msis(phb);
 
         */
        ppc_md.pcibios_fixup = pnv_pci_ioda_fixup;
 
-       if (phb->type == PNV_PHB_NPU)
+       if (phb->type == PNV_PHB_NPU) {
                hose->controller_ops = pnv_npu_ioda_controller_ops;
-       else
+       } else {
+               phb->dma_dev_setup = pnv_pci_ioda_dma_dev_setup;
                hose->controller_ops = pnv_pci_ioda_controller_ops;
+       }
 
 #ifdef CONFIG_PCI_IOV
        ppc_md.pcibios_fixup_sriov = pnv_pci_ioda_fixup_iov_resources;