diff --git a/drivers/gpu/nvgpu/os/linux/linux-dma.c b/drivers/gpu/nvgpu/os/linux/linux-dma.c
index 4687f0d72..83e238d83 100644
--- a/drivers/gpu/nvgpu/os/linux/linux-dma.c
+++ b/drivers/gpu/nvgpu/os/linux/linux-dma.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -16,7 +16,7 @@
 
 #include <linux/dma-mapping.h>
 #include <linux/slab.h>
-#include <linux/version.h>
+#include <linux/iommu.h>
 
 #include <nvgpu/log.h>
 #include <nvgpu/dma.h>
@@ -604,7 +604,7 @@ bool nvgpu_iommuable(struct gk20a *g)
 	 * Check against the nvgpu device to see if it's been marked as
 	 * IOMMU'able.
 	 */
-	if (dev->archdata.iommu == NULL)
+	if (iommu_get_domain_for_dev(dev) == NULL)
 		return false;
 #endif
 
diff --git a/drivers/gpu/nvgpu/os/linux/platform_gp10b_tegra.c b/drivers/gpu/nvgpu/os/linux/platform_gp10b_tegra.c
index f356ab9ff..e4c31979d 100644
--- a/drivers/gpu/nvgpu/os/linux/platform_gp10b_tegra.c
+++ b/drivers/gpu/nvgpu/os/linux/platform_gp10b_tegra.c
@@ -19,6 +19,7 @@
 #include <linux/debugfs.h>
 #include <linux/dma-buf.h>
 #include <linux/reset.h>
+#include <linux/iommu.h>
 #ifdef CONFIG_TEGRA_BWMGR
 #include <linux/platform/tegra/emc_bwmgr.h>
 #endif
@@ -204,7 +205,8 @@ static int gp10b_tegra_probe(struct device *dev)
 	if (ret)
 		return ret;
 
-	platform->disable_bigpage = !dev->archdata.iommu && (PAGE_SIZE < SZ_64K);
+	platform->disable_bigpage = !iommu_get_domain_for_dev(dev) &&
+					(PAGE_SIZE < SZ_64K);
 
 #ifdef CONFIG_OF
 	of_chosen = of_find_node_by_path("/chosen");
diff --git a/drivers/gpu/nvgpu/os/linux/platform_gv11b_tegra.c b/drivers/gpu/nvgpu/os/linux/platform_gv11b_tegra.c
index 6e8477934..6311c9f10 100644
--- a/drivers/gpu/nvgpu/os/linux/platform_gv11b_tegra.c
+++ b/drivers/gpu/nvgpu/os/linux/platform_gv11b_tegra.c
@@ -22,6 +22,7 @@
 #include <linux/reset.h>
 #include <linux/hashtable.h>
 #include <linux/clk.h>
+#include <linux/iommu.h>
 #ifdef CONFIG_TEGRA_BWMGR
 #include <linux/platform/tegra/emc_bwmgr.h>
 #endif
@@ -94,7 +95,8 @@ static int gv11b_tegra_probe(struct device *dev)
 	if (err)
 		return err;
 
-	platform->disable_bigpage = !dev->archdata.iommu && (PAGE_SIZE < SZ_64K);
+	platform->disable_bigpage = !(iommu_get_domain_for_dev(dev)) &&
+					(PAGE_SIZE < SZ_64K);
 
 #ifdef CONFIG_OF
 	of_chosen = of_find_node_by_path("/chosen");