diff --git a/kernel/common/inc/nv-linux.h b/kernel/common/inc/nv-linux.h
index 4043f37..56eb2cb 100644
--- a/kernel/common/inc/nv-linux.h
+++ b/kernel/common/inc/nv-linux.h
@@ -498,11 +498,17 @@ extern int nv_pat_mode;
 /* keep track of memory usage */
 #include "nv-memdbg.h"
 
 static inline void *nv_vmalloc(unsigned long size)
 {
+#if NV_VMALLOC_ARGUMENT_COUNT == 3
     void *ptr = __vmalloc(size, GFP_KERNEL, PAGE_KERNEL);
+#elif NV_VMALLOC_ARGUMENT_COUNT == 2
+    void *ptr = __vmalloc(size, GFP_KERNEL);
+#else
+#error "NV_VMALLOC_ARGUMENT_COUNT value unrecognized!"
+#endif
     if (ptr)
         NV_MEMDBG_ADD(ptr, size);
     return ptr;
 }
 
diff --git a/kernel/common/inc/nv-mm.h b/kernel/common/inc/nv-mm.h
index 88c33fe..4a88abe 100644
--- a/kernel/common/inc/nv-mm.h
+++ b/kernel/common/inc/nv-mm.h
@@ -23,10 +23,14 @@
 #ifndef __NV_MM_H__
 #define __NV_MM_H__
 
 #include "conftest.h"
 
+#if defined(NV_MM_STRUCT_HAS_MMAP_LOCK)
+#define mmap_sem mmap_lock
+#endif
+
 #ifndef NV_VMF_INSERT_PFN_PRESENT
 typedef int vm_fault_t;
 #endif
 
 /* get_user_pages
diff --git a/kernel/conftest.sh b/kernel/conftest.sh
index dcd6dd5..1be6685 100755
--- a/kernel/conftest.sh
+++ b/kernel/conftest.sh
@@ -2032,10 +2032,53 @@ compile_test() {
             }"
 
             compile_check_conftest "$CODE" "NV_VZALLOC_PRESENT" "" "functions"
         ;;
 
+        vmalloc_argument_count)
+            #
+            # Determine how many arguments __vmalloc takes.
+            #
+            # Changed by commit fc3af83c4fca ("mm: remove the pgprot argument
+            # to __vmalloc")
+            #
+            echo "$CONFTEST_PREAMBLE
+            #include <linux/mm.h>
+            #include <linux/vmalloc.h>
+            void conftest_vmalloc_argument_count(void) {
+                __vmalloc(0, GFP_KERNEL, PAGE_KERNEL);
+            }" > conftest$$.c
+
+            $CC $CFLAGS -c conftest$$.c > /dev/null 2>&1
+            rm -f conftest$$.c
+
+            if [ -f conftest$$.o ]; then
+                echo "#define NV_VMALLOC_ARGUMENT_COUNT 3" | append_conftest "functions"
+            else
+                echo "#define NV_VMALLOC_ARGUMENT_COUNT 2" | append_conftest "functions"
+            fi
+
+            rm -f conftest$$.o
+        ;;
+
+        mm_struct_has_mmap_lock)
+            #
+            # Determine if the mm_struct structure has 'mmap_lock'.
+            #
+            # Changed by commit ea7b54944ef9 ("mmap locking API: rename mmap_sem
+            # to mmap_lock")
+            #
+            CODE="
+            #include <linux/mm.h>
+
+            int conftest_mm_struct_has_mmap_lock(void) {
+                return offsetof(struct mm_struct, mmap_lock);
+            }"
+
+            compile_check_conftest "$CODE" "NV_MM_STRUCT_HAS_MMAP_LOCK" "" "types"
+        ;;
+
         drm_driver_has_set_busid)
             #
             # Determine if the drm_driver structure has a 'set_busid' callback
             # field.
             #
diff --git a/kernel/nvidia/nvidia.Kbuild b/kernel/nvidia/nvidia.Kbuild
index a565f86..b5fea30 100644
--- a/kernel/nvidia/nvidia.Kbuild
+++ b/kernel/nvidia/nvidia.Kbuild
@@ -141,10 +141,11 @@ NV_CONFTEST_FUNCTION_COMPILE_TESTS += kernel_write
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += kthread_create_on_node
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += of_find_matching_node
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += dev_is_pci
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += dma_direct_map_resource
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += tegra_get_platform
+NV_CONFTEST_FUNCTION_COMPILE_TESTS += vmalloc_argument_count
 
 NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_gpl_of_node_to_nid
 NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_swiotlb_map_sg_attrs
 NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_swiotlb_dma_ops
 
@@ -164,10 +165,11 @@ NV_CONFTEST_TYPE_COMPILE_TESTS += address_space_init_once
 NV_CONFTEST_TYPE_COMPILE_TESTS += vm_ops_fault_removed_vma_arg
 NV_CONFTEST_TYPE_COMPILE_TESTS += vmbus_channel_has_ringbuffer_page
 NV_CONFTEST_TYPE_COMPILE_TESTS += device_driver_of_match_table
 NV_CONFTEST_TYPE_COMPILE_TESTS += device_of_node
 NV_CONFTEST_TYPE_COMPILE_TESTS += node_states_n_memory
+NV_CONFTEST_TYPE_COMPILE_TESTS += mm_struct_has_mmap_lock
 
 NV_CONFTEST_GENERIC_COMPILE_TESTS += dom0_kernel_present
 NV_CONFTEST_GENERIC_COMPILE_TESTS += nvidia_vgpu_hyperv_available
 NV_CONFTEST_GENERIC_COMPILE_TESTS += nvidia_vgpu_kvm_build
 NV_CONFTEST_GENERIC_COMPILE_TESTS += nvidia_grid_build
-- 
2.28.0

