diff --git a/drivers/misc/nvscic2c-pcie/pci-client.c b/drivers/misc/nvscic2c-pcie/pci-client.c
index 2eb0ad51..15cfe6d8 100644
--- a/drivers/misc/nvscic2c-pcie/pci-client.c
+++ b/drivers/misc/nvscic2c-pcie/pci-client.c
@@ -3,6 +3,8 @@
 
 #define pr_fmt(fmt)	"nvscic2c-pcie: pci-client: " fmt
 
+#include <nvidia/conftest.h>
+
 #include <linux/dma-buf.h>
 #include <linux/dma-map-ops.h>
 #include <linux/errno.h>
@@ -246,7 +248,12 @@ allocate_edma_rx_desc_iova(struct pci_client_t *ctx)
 		goto err;
 	}
 	prot = (IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
-	ret = iommu_map(ctx->domain, ctx->edma_ch_desc_iova, phys_addr, EDMA_CH_DESC_SZ, prot);
+	ret = iommu_map(ctx->domain, ctx->edma_ch_desc_iova, phys_addr,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+		EDMA_CH_DESC_SZ, prot, GFP_KERNEL);
+#else
+		EDMA_CH_DESC_SZ, prot);
+#endif
 	if (ret) {
 		pr_err("pci client failed to map iova to 60K physical backing\n");
 		goto err;
@@ -408,7 +415,12 @@ pci_client_map_addr(void *pci_client_h, u64 to_iova, phys_addr_t paddr,
 	if (WARN_ON(!ctx || !to_iova || !paddr || !size))
 		return -EINVAL;
 
-	return iommu_map(ctx->domain, to_iova, paddr, size, prot);
+	return iommu_map(ctx->domain, to_iova, paddr, size,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+		prot, GFP_KERNEL);
+#else
+		prot);
+#endif
 }
 
 size_t
diff --git a/drivers/pci/endpoint/functions/pci-epf-tegra-vnet.c b/drivers/pci/endpoint/functions/pci-epf-tegra-vnet.c
index 5d4d50be..61cd8c11 100644
--- a/drivers/pci/endpoint/functions/pci-epf-tegra-vnet.c
+++ b/drivers/pci/endpoint/functions/pci-epf-tegra-vnet.c
@@ -321,7 +321,11 @@ static void tvnet_ep_alloc_empty_buffers(struct pci_epf_tvnet *tvnet)
 		}
 
 		ret = iommu_map(domain, iova, page_to_phys(page), PAGE_SIZE,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+				IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
+#else
 				IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
+#endif
 		if (ret < 0) {
 			dev_err(tvnet->fdev, "%s: iommu_map(RAM) failed: %d\n",
 				__func__, ret);
@@ -1237,7 +1241,11 @@ static int tvnet_ep_pci_epf_setup_irqsp(struct pci_epf_tvnet *tvnet)
 	syncpt_addr = nvhost_interrupt_syncpt_get_syncpt_addr(ctrl_irqsp->is);
 #endif
 	ret = iommu_map(domain, amap->iova, syncpt_addr, PAGE_SIZE,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
+#else
 			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
+#endif
 	if (ret < 0) {
 		dev_err(fdev, "%s: iommu_map of ctrlsp mem failed: %d\n",
 			__func__, ret);
@@ -1256,7 +1264,11 @@ static int tvnet_ep_pci_epf_setup_irqsp(struct pci_epf_tvnet *tvnet)
 	syncpt_addr = nvhost_interrupt_syncpt_get_syncpt_addr(data_irqsp->is);
 #endif
 	ret = iommu_map(domain, amap->iova + PAGE_SIZE, syncpt_addr, PAGE_SIZE,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
+#else
 			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
+#endif
 	if (ret < 0) {
 		dev_err(fdev, "%s: iommu_map of datasp mem failed: %d\n",
 			__func__, ret);
@@ -1323,7 +1335,11 @@ static int tvnet_ep_alloc_single_page_bar0_mem(struct pci_epf *epf,
 	}
 
 	ret = iommu_map(domain, amap->iova, page_to_phys(amap->page), PAGE_SIZE,
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
+#else
 			IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
+#endif
 	if (ret < 0) {
 		dev_err(tvnet->fdev, "%s: type: %d iommu_map(RAM) failed: %d\n",
 			__func__, type, ret);
@@ -1404,7 +1420,11 @@ static int tvnet_ep_alloc_multi_page_bar0_mem(struct pci_epf *epf,
 	}
 
 	ret = iommu_map(domain, amap->iova, page_to_phys(amap->page),
+#if defined(NV_IOMMU_MAP_HAS_GFP_ARG)
+			amap->size, IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE, GFP_KERNEL);
+#else
 			amap->size, IOMMU_CACHE | IOMMU_READ | IOMMU_WRITE);
+#endif
 	if (ret < 0) {
 		dev_err(tvnet->fdev, "%s: iommu_map(RAM) failed: %d\n",
 			__func__, ret);