diff --git a/os/linux/mem_mgmt_utils.c b/os/linux/mem_mgmt_utils.c
index 145242b..065355f 100644
--- a/os/linux/mem_mgmt_utils.c
+++ b/os/linux/mem_mgmt_utils.c
@@ -128,28 +128,26 @@ static int tegra_hwpm_dma_map_mem_bytes_buffer(struct tegra_soc_hwpm *hwpm,
 		sg_dma_address(hwpm->mem_mgmt->mem_bytes_sgt->sgl);
 
 #if defined(CONFIG_TEGRA_HWPM_OOT)
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 18, 0)
+	hwpm->mem_mgmt->mem_bytes_kernel =
+		(struct iosys_map *) &hwpm->mem_mgmt->mem_bytes_map;
+	err = dma_buf_vmap(hwpm->mem_mgmt->mem_bytes_dma_buf,
+		(struct iosys_map *)hwpm->mem_mgmt->mem_bytes_kernel);
+#else
 	hwpm->mem_mgmt->mem_bytes_kernel =
 		(struct dma_buf_map *) &hwpm->mem_mgmt->mem_bytes_map;
 	err = dma_buf_vmap(hwpm->mem_mgmt->mem_bytes_dma_buf,
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 18, 0)
-		(struct iosys_map *)hwpm->mem_mgmt->mem_bytes_kernel);
-#else
 		(struct dma_buf_map *)hwpm->mem_mgmt->mem_bytes_kernel);
 #endif
-	if (err != 0) {
-		tegra_hwpm_err(hwpm,
-			"Unable to map mem_bytes buffer into kernel VA space");
-		return -ENOMEM;
-	}
-#else
+#else /* !CONFIG_TEGRA_HWPM_OOT */
 	hwpm->mem_mgmt->mem_bytes_kernel =
 		dma_buf_vmap(hwpm->mem_mgmt->mem_bytes_dma_buf);
+#endif
 	if (!hwpm->mem_mgmt->mem_bytes_kernel) {
 		tegra_hwpm_err(hwpm,
 			"Unable to map mem_bytes buffer into kernel VA space");
 		return -ENOMEM;
 	}
-#endif
 
 	memset(hwpm->mem_mgmt->mem_bytes_kernel, 0, 32);
 
diff --git a/os/linux/mem_mgmt_utils.h b/os/linux/mem_mgmt_utils.h
index dd730b5..c6b63da 100644
--- a/os/linux/mem_mgmt_utils.h
+++ b/os/linux/mem_mgmt_utils.h
@@ -16,8 +16,13 @@
 
 #include <linux/types.h>
 #if defined(CONFIG_TEGRA_HWPM_OOT)
+#include <linux/version.h>
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 18, 0)
+#include <linux/iosys-map.h>
+#else
 #include <linux/dma-buf-map.h>
 #endif
+#endif
 
 /* This macro is copy of TEGRA_SOC_HWPM_MEM_BYTES_INVALID */
 #define TEGRA_HWPM_MEM_BYTES_INVALID	0xffffffff
@@ -42,8 +47,12 @@ struct tegra_hwpm_mem_mgmt {
 	u64 mem_bytes_buf_va;
 	void *mem_bytes_kernel;
 #if defined(CONFIG_TEGRA_HWPM_OOT)
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 18, 0)
+	struct iosys_map mem_bytes_map;
+#else
 	struct dma_buf_map mem_bytes_map;
 #endif
+#endif
 };
 
 struct tegra_hwpm_allowlist_map {