535.288.01

CHANGELOG.md

@@ -1,154 +0,0 @@
-# Changelog
-
-## Release 535 Entries
-
-### [535.86.05] 2023-07-18
-
-### [535.54.03] 2023-06-14
-
-### [535.43.02] 2023-05-30
-
-#### Fixed
-
-- Fixed console restore with traditional VGA consoles.
-
-#### Added
-
-- Added support for Run Time D3 (RTD3) on Ampere and later GPUs.
-- Added support for G-Sync on desktop GPUs.
-
-## Release 530 Entries
-
-### [530.41.03] 2023-03-23
-
-### [530.30.02] 2023-02-28
-
-#### Changed
-
-- GSP firmware is now distributed as `gsp_tu10x.bin` and `gsp_ga10x.bin` to better reflect the GPU architectures supported by each firmware file in this release.
-    - The .run installer will continue to install firmware to /lib/firmware/nvidia/<version> and the nvidia.ko kernel module will load the appropriate firmware for each GPU at runtime.
-  
-#### Fixed
-
-- Add support for resizable BAR on Linux when NVreg_EnableResizableBar=1 module param is set. [#3](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/3) by @sjkelly
-
-#### Added
-
-- Support for power management features like Suspend, Hibernate and Resume.
-
-## Release 525 Entries
-
-### [525.116.04] 2023-05-09
-
-### [525.116.03] 2023-04-25
-
-### [525.105.17] 2023-03-30
-
-### [525.89.02] 2023-02-08
-
-### [525.85.12] 2023-01-30
-
-### [525.85.05] 2023-01-19
-
-#### Fixed
-
-- Fix build problems with Clang 15.0, [#377](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/377) by @ptr1337
-
-### [525.78.01] 2023-01-05
-
-### [525.60.13] 2022-12-05
-
-### [525.60.11] 2022-11-28
-
-#### Fixed
-
-- Fixed nvenc compatibility with usermode clients [#104](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/104)
-
-### [525.53] 2022-11-10
-
-#### Changed
-
-- GSP firmware is now distributed as multiple firmware files: this release has `gsp_tu10x.bin` and `gsp_ad10x.bin` replacing `gsp.bin` from previous releases.
-    - Each file is named after a GPU architecture and supports GPUs from one or more architectures. This allows GSP firmware to better leverage each architecture's capabilities.
-    - The .run installer will continue to install firmware to `/lib/firmware/nvidia/<version>` and the `nvidia.ko` kernel module will load the appropriate firmware for each GPU at runtime.
-
-#### Fixed
-
-- Add support for IBT (indirect branch tracking) on supported platforms, [#256](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/256) by @rnd-ash
-- Return EINVAL when [failing to] allocating memory, [#280](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/280) by @YusufKhan-gamedev
-- Fix various typos in nvidia/src/kernel, [#16](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/16) by @alexisgeoffrey
-- Added support for rotation in X11, Quadro Sync, Stereo, and YUV 4:2:0 on Turing.
-
-## Release 520 Entries
-
-### [520.61.07] 2022-10-20
-
-### [520.56.06] 2022-10-12
-
-#### Added
-
-- Introduce support for GeForce RTX 4090 GPUs.
-
-### [520.61.05] 2022-10-10
-
-#### Added
-
-- Introduce support for NVIDIA H100 GPUs.
-
-#### Fixed
-
-- Fix/Improve Makefile, [#308](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/308/) by @izenynn
-- Make nvLogBase2 more efficient, [#177](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/177/) by @DMaroo
-- nv-pci: fixed always true expression, [#195](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/195/) by @ValZapod
-
-## Release 515 Entries
-
-### [515.76] 2022-09-20
-
-#### Fixed
-
-- Improved compatibility with new Linux kernel releases
-- Fixed possible excessive GPU power draw on an idle X11 or Wayland desktop when driving high resolutions or refresh rates
-
-### [515.65.07] 2022-10-19
-
-### [515.65.01] 2022-08-02
-
-#### Fixed
-
-- Collection of minor fixes to issues, [#6](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/61) by @Joshua-Ashton
-- Remove unnecessary use of acpi_bus_get_device().
-
-### [515.57] 2022-06-28
-
-#### Fixed
-
-- Backtick is deprecated, [#273](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/273) by @arch-user-france1
-
-### [515.48.07] 2022-05-31
-
-#### Added
-
-- List of compatible GPUs in README.md.
-
-#### Fixed
-
-- Fix various README capitalizations, [#8](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/8) by @27lx 
-- Automatically tag bug report issues, [#15](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/15) by @thebeanogamer
-- Improve conftest.sh Script, [#37](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/37) by @Nitepone
-- Update HTTP link to HTTPS, [#101](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/101) by @alcaparra
-- moved array sanity check to before the array access, [#117](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/117) by @RealAstolfo
-- Fixed some typos, [#122](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/122) by @FEDOyt
-- Fixed capitalization, [#123](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/123) by @keroeslux
-- Fix typos in NVDEC Engine Descriptor, [#126](https://github.com/NVIDIA/open-gpu-kernel-modules/pull/126) from @TrickyDmitriy
-- Extranous apostrohpes in a makefile script [sic], [#14](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/14) by @kiroma
-- HDMI no audio @ 4K above 60Hz, [#75](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/75) by @adolfotregosa
-- dp_configcaps.cpp:405: array index sanity check in wrong place?, [#110](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/110) by @dcb314
-- NVRM kgspInitRm_IMPL: missing NVDEC0 engine, cannot initialize GSP-RM, [#116](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/116) by @kfazz
-- ERROR: modpost: "backlight_device_register" [...nvidia-modeset.ko] undefined, [#135](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/135) by @sndirsch
-- aarch64 build fails, [#151](https://github.com/NVIDIA/open-gpu-kernel-modules/issues/151) by @frezbo
-
-### [515.43.04] 2022-05-11
-
-- Initial release.
-

README.md

@@ -1,7 +1,7 @@
 # NVIDIA Linux Open GPU Kernel Module Source
 
 This is the source release of the NVIDIA Linux open GPU kernel modules,
-version 535.86.05.
+version 535.288.01.
 
 
 ## How to Build
@@ -17,7 +17,7 @@ as root:
 
 Note that the kernel modules built here must be used with GSP
 firmware and user-space NVIDIA GPU driver components from a corresponding
-535.86.05 driver release.  This can be achieved by installing
+535.288.01 driver release.  This can be achieved by installing
 the NVIDIA GPU driver from the .run file using the `--no-kernel-modules`
 option.  E.g.,
 
@@ -180,7 +180,7 @@ software applications.
 ## Compatible GPUs
 
 The open-gpu-kernel-modules can be used on any Turing or later GPU
-(see the table below). However, in the 535.86.05 release,
+(see the table below). However, in the 535.288.01 release,
 GeForce and Workstation support is still considered alpha-quality.
 
 To enable use of the open kernel modules on GeForce and Workstation GPUs,
@@ -188,7 +188,7 @@ set the "NVreg_OpenRmEnableUnsupportedGpus" nvidia.ko kernel module
 parameter to 1. For more details, see the NVIDIA GPU driver end user
 README here:
 
-https://us.download.nvidia.com/XFree86/Linux-x86_64/535.86.05/README/kernel_open.html
+https://us.download.nvidia.com/XFree86/Linux-x86_64/535.288.01/README/kernel_open.html
 
 In the below table, if three IDs are listed, the first is the PCI Device 
 ID, the second is the PCI Subsystem Vendor ID, and the third is the PCI
@@ -648,6 +648,7 @@ Subsystem Device ID.
 | NVIDIA T1000 8GB                                | 1FF0 17AA 1612 |
 | NVIDIA T400 4GB                                 | 1FF2 1028 1613 |
 | NVIDIA T400 4GB                                 | 1FF2 103C 1613 |
+| NVIDIA T400E                                    | 1FF2 103C 18FF |
 | NVIDIA T400 4GB                                 | 1FF2 103C 8A80 |
 | NVIDIA T400 4GB                                 | 1FF2 10DE 1613 |
 | NVIDIA T400 4GB                                 | 1FF2 17AA 1613 |
@@ -658,6 +659,7 @@ Subsystem Device ID.
 | NVIDIA A100-SXM4-80GB                           | 20B2 10DE 147F |
 | NVIDIA A100-SXM4-80GB                           | 20B2 10DE 1622 |
 | NVIDIA A100-SXM4-80GB                           | 20B2 10DE 1623 |
+| NVIDIA PG509-210                                | 20B2 10DE 1625 |
 | NVIDIA A100-SXM-64GB                            | 20B3 10DE 14A7 |
 | NVIDIA A100-SXM-64GB                            | 20B3 10DE 14A8 |
 | NVIDIA A100 80GB PCIe                           | 20B5 10DE 1533 |
@@ -665,6 +667,8 @@ Subsystem Device ID.
 | NVIDIA PG506-232                                | 20B6 10DE 1492 |
 | NVIDIA A30                                      | 20B7 10DE 1532 |
 | NVIDIA A30                                      | 20B7 10DE 1804 |
+| NVIDIA A30                                      | 20B7 10DE 1852 |
+| NVIDIA A800-SXM4-40GB                           | 20BD 10DE 17F4 |
 | NVIDIA A100-PCIE-40GB                           | 20F1 10DE 145F |
 | NVIDIA A800-SXM4-80GB                           | 20F3 10DE 179B |
 | NVIDIA A800-SXM4-80GB                           | 20F3 10DE 179C |
@@ -676,6 +680,11 @@ Subsystem Device ID.
 | NVIDIA A800-SXM4-80GB                           | 20F3 10DE 17A2 |
 | NVIDIA A800 80GB PCIe                           | 20F5 10DE 1799 |
 | NVIDIA A800 80GB PCIe LC                        | 20F5 10DE 179A |
+| NVIDIA A800 40GB Active                         | 20F6 1028 180A |
+| NVIDIA A800 40GB Active                         | 20F6 103C 180A |
+| NVIDIA A800 40GB Active                         | 20F6 10DE 180A |
+| NVIDIA A800 40GB Active                         | 20F6 17AA 180A |
+| NVIDIA AX800                                    | 20FD 10DE 17F8 |
 | NVIDIA GeForce GTX 1660 Ti                      | 2182           |
 | NVIDIA GeForce GTX 1660                         | 2184           |
 | NVIDIA GeForce GTX 1650 SUPER                   | 2187           |
@@ -734,13 +743,21 @@ Subsystem Device ID.
 | NVIDIA A10                                      | 2236 10DE 1482 |
 | NVIDIA A10G                                     | 2237 10DE 152F |
 | NVIDIA A10M                                     | 2238 10DE 1677 |
+| NVIDIA H100 NVL                                 | 2321 10DE 1839 |
 | NVIDIA H800 PCIe                                | 2322 10DE 17A4 |
 | NVIDIA H800                                     | 2324 10DE 17A6 |
 | NVIDIA H800                                     | 2324 10DE 17A8 |
+| NVIDIA H20                                      | 2329 10DE 198B |
+| NVIDIA H20                                      | 2329 10DE 198C |
+| NVIDIA H20-3e                                   | 232C 10DE 2063 |
 | NVIDIA H100 80GB HBM3                           | 2330 10DE 16C0 |
 | NVIDIA H100 80GB HBM3                           | 2330 10DE 16C1 |
 | NVIDIA H100 PCIe                                | 2331 10DE 1626 |
 | NVIDIA H100                                     | 2339 10DE 17FC |
+| NVIDIA H800 NVL                                 | 233A 10DE 183A |
+| NVIDIA GH200 120GB                              | 2342 10DE 16EB |
+| NVIDIA GH200 120GB                              | 2342 10DE 1805 |
+| NVIDIA GH200 480GB                              | 2342 10DE 1809 |
 | NVIDIA GeForce RTX 3060 Ti                      | 2414           |
 | NVIDIA GeForce RTX 3080 Ti Laptop GPU           | 2420           |
 | NVIDIA RTX A5500 Laptop GPU                     | 2438           |
@@ -793,6 +810,7 @@ Subsystem Device ID.
 | NVIDIA RTX A2000 12GB                           | 2571 10DE 1611 |
 | NVIDIA RTX A2000 12GB                           | 2571 17AA 1611 |
 | NVIDIA GeForce RTX 3050                         | 2582           |
+| NVIDIA GeForce RTX 3050                         | 2584           |
 | NVIDIA GeForce RTX 3050 Ti Laptop GPU           | 25A0           |
 | NVIDIA GeForce RTX 3050Ti Laptop GPU            | 25A0 103C 8928 |
 | NVIDIA GeForce RTX 3050Ti Laptop GPU            | 25A0 103C 89F9 |
@@ -829,12 +847,22 @@ Subsystem Device ID.
 | NVIDIA RTX 6000 Ada Generation                  | 26B1 103C 16A1 |
 | NVIDIA RTX 6000 Ada Generation                  | 26B1 10DE 16A1 |
 | NVIDIA RTX 6000 Ada Generation                  | 26B1 17AA 16A1 |
+| NVIDIA RTX 5000 Ada Generation                  | 26B2 1028 17FA |
+| NVIDIA RTX 5000 Ada Generation                  | 26B2 103C 17FA |
+| NVIDIA RTX 5000 Ada Generation                  | 26B2 10DE 17FA |
+| NVIDIA RTX 5000 Ada Generation                  | 26B2 17AA 17FA |
+| NVIDIA RTX 5880 Ada Generation                  | 26B3 103C 1934 |
+| NVIDIA RTX 5880 Ada Generation                  | 26B3 10DE 1934 |
 | NVIDIA L40                                      | 26B5 10DE 169D |
 | NVIDIA L40                                      | 26B5 10DE 17DA |
+| NVIDIA L40S                                     | 26B9 10DE 1851 |
+| NVIDIA L40S                                     | 26B9 10DE 18CF |
+| NVIDIA L20                                      | 26BA 10DE 1957 |
 | NVIDIA GeForce RTX 4080                         | 2704           |
 | NVIDIA GeForce RTX 4090 Laptop GPU              | 2717           |
 | NVIDIA RTX 5000 Ada Generation Laptop GPU       | 2730           |
 | NVIDIA GeForce RTX 4090 Laptop GPU              | 2757           |
+| NVIDIA RTX 5000 Ada Generation Embedded GPU     | 2770           |
 | NVIDIA GeForce RTX 4070 Ti                      | 2782           |
 | NVIDIA GeForce RTX 4070                         | 2786           |
 | NVIDIA GeForce RTX 4080 Laptop GPU              | 27A0           |
@@ -842,11 +870,21 @@ Subsystem Device ID.
 | NVIDIA RTX 4000 SFF Ada Generation              | 27B0 103C 16FA |
 | NVIDIA RTX 4000 SFF Ada Generation              | 27B0 10DE 16FA |
 | NVIDIA RTX 4000 SFF Ada Generation              | 27B0 17AA 16FA |
+| NVIDIA RTX 4500 Ada Generation                  | 27B1 1028 180C |
+| NVIDIA RTX 4500 Ada Generation                  | 27B1 103C 180C |
+| NVIDIA RTX 4500 Ada Generation                  | 27B1 10DE 180C |
+| NVIDIA RTX 4500 Ada Generation                  | 27B1 17AA 180C |
+| NVIDIA RTX 4000 Ada Generation                  | 27B2 1028 181B |
+| NVIDIA RTX 4000 Ada Generation                  | 27B2 103C 181B |
+| NVIDIA RTX 4000 Ada Generation                  | 27B2 10DE 181B |
+| NVIDIA RTX 4000 Ada Generation                  | 27B2 17AA 181B |
+| NVIDIA L2                                       | 27B6 10DE 1933 |
 | NVIDIA L4                                       | 27B8 10DE 16CA |
 | NVIDIA L4                                       | 27B8 10DE 16EE |
 | NVIDIA RTX 4000 Ada Generation Laptop GPU       | 27BA           |
 | NVIDIA RTX 3500 Ada Generation Laptop GPU       | 27BB           |
 | NVIDIA GeForce RTX 4080 Laptop GPU              | 27E0           |
+| NVIDIA RTX 3500 Ada Generation Embedded GPU     | 27FB           |
 | NVIDIA GeForce RTX 4060 Ti                      | 2803           |
 | NVIDIA GeForce RTX 4060 Ti                      | 2805           |
 | NVIDIA GeForce RTX 4070 Laptop GPU              | 2820           |
@@ -856,5 +894,9 @@ Subsystem Device ID.
 | NVIDIA GeForce RTX 4060 Laptop GPU              | 28A0           |
 | NVIDIA GeForce RTX 4050 Laptop GPU              | 28A1           |
 | NVIDIA RTX 2000 Ada Generation Laptop GPU       | 28B8           |
+| NVIDIA RTX 1000 Ada Generation Laptop GPU       | 28B9           |
+| NVIDIA RTX 500 Ada Generation Laptop GPU        | 28BA           |
+| NVIDIA RTX 500 Ada Generation Laptop GPU        | 28BB           |
 | NVIDIA GeForce RTX 4060 Laptop GPU              | 28E0           |
 | NVIDIA GeForce RTX 4050 Laptop GPU              | 28E1           |
+| NVIDIA RTX 2000 Ada Generation Embedded GPU     | 28F8           |

kernel-open/Kbuild

@@ -57,72 +57,82 @@ ifeq ($(NV_UNDEF_BEHAVIOR_SANITIZER),1)
  UBSAN_SANITIZE := y
 endif
 
+#
+# Command to create a symbolic link, explicitly resolving the symlink target
+# to an absolute path to abstract away the difference between Linux < 6.13,
+# where the CWD is the Linux kernel source tree for Kbuild extmod builds, and
+# Linux >= 6.13, where the CWD is the external module source tree.
+#
+# This is used to create the nv*-kernel.o -> nv*-kernel.o_binary symlinks for
+# kernel modules which use precompiled binary object files.
+#
+
+quiet_cmd_symlink = SYMLINK $@
+ cmd_symlink = ln -sf $(abspath $<) $@
+
+
 $(foreach _module, $(NV_KERNEL_MODULES), \
  $(eval include $(src)/$(_module)/$(_module).Kbuild))
 
 
-#
-# Define CFLAGS that apply to all the NVIDIA kernel modules. EXTRA_CFLAGS
-# is deprecated since 2.6.24 in favor of ccflags-y, but we need to support
-# older kernels which do not have ccflags-y. Newer kernels append
-# $(EXTRA_CFLAGS) to ccflags-y for compatibility.
-#
-
-EXTRA_CFLAGS += -I$(src)/common/inc
-EXTRA_CFLAGS += -I$(src)
-EXTRA_CFLAGS += -Wall $(DEFINES) $(INCLUDES) -Wno-cast-qual -Wno-error -Wno-format-extra-args
-EXTRA_CFLAGS += -D__KERNEL__ -DMODULE -DNVRM
-EXTRA_CFLAGS += -DNV_VERSION_STRING=\"535.86.05\"
+ccflags-y += -I$(src)/common/inc
+ccflags-y += -I$(src)
+ccflags-y += -Wall $(DEFINES) $(INCLUDES) -Wno-cast-qual -Wno-format-extra-args
+ccflags-y += -D__KERNEL__ -DMODULE -DNVRM
+ccflags-y += -DNV_VERSION_STRING=\"535.288.01\"
 
 ifneq ($(SYSSRCHOST1X),)
- EXTRA_CFLAGS += -I$(SYSSRCHOST1X)
+ ccflags-y += -I$(SYSSRCHOST1X)
 endif
 
-EXTRA_CFLAGS += -Wno-unused-function
+ccflags-y += -Wno-unused-function
 
 ifneq ($(NV_BUILD_TYPE),debug)
- EXTRA_CFLAGS += -Wuninitialized
+ ccflags-y += -Wuninitialized
 endif
 
-EXTRA_CFLAGS += -fno-strict-aliasing
+ccflags-y += -fno-strict-aliasing
 
 ifeq ($(ARCH),arm64)
- EXTRA_CFLAGS += -mstrict-align
+ ccflags-y += -mstrict-align
 endif
 
 ifeq ($(NV_BUILD_TYPE),debug)
- EXTRA_CFLAGS += -g
- EXTRA_CFLAGS += $(call cc-option,-gsplit-dwarf,)
+ ccflags-y += -g
+ ccflags-y += $(call cc-option,-gsplit-dwarf,)
 endif
 
-EXTRA_CFLAGS += -ffreestanding
+ccflags-y += -ffreestanding
 
 ifeq ($(ARCH),arm64)
- EXTRA_CFLAGS += -mgeneral-regs-only -march=armv8-a
- EXTRA_CFLAGS += $(call cc-option,-mno-outline-atomics,)
+ ccflags-y += -mgeneral-regs-only -march=armv8-a
+ ccflags-y += $(call cc-option,-mno-outline-atomics,)
 endif
 
 ifeq ($(ARCH),x86_64)
- EXTRA_CFLAGS += -mno-red-zone -mcmodel=kernel
+ ccflags-y += -mno-red-zone -mcmodel=kernel
 endif
 
 ifeq ($(ARCH),powerpc)
- EXTRA_CFLAGS += -mlittle-endian -mno-strict-align -mno-altivec
+ ccflags-y += -mlittle-endian -mno-strict-align -mno-altivec
 endif
 
-EXTRA_CFLAGS += -DNV_UVM_ENABLE
-EXTRA_CFLAGS += $(call cc-option,-Werror=undef,)
-EXTRA_CFLAGS += -DNV_SPECTRE_V2=$(NV_SPECTRE_V2)
-EXTRA_CFLAGS += -DNV_KERNEL_INTERFACE_LAYER
+ccflags-y += -DNV_UVM_ENABLE
+ccflags-y += $(call cc-option,-Werror=undef,)
+ccflags-y += -DNV_SPECTRE_V2=$(NV_SPECTRE_V2)
+ccflags-y += -DNV_KERNEL_INTERFACE_LAYER
 
 #
 # Detect SGI UV systems and apply system-specific optimizations.
 #
 
 ifneq ($(wildcard /proc/sgi_uv),)
- EXTRA_CFLAGS += -DNV_CONFIG_X86_UV
+ ccflags-y += -DNV_CONFIG_X86_UV
 endif
 
+ifdef VGX_FORCE_VFIO_PCI_CORE
+ ccflags-y += -DNV_VGPU_FORCE_VFIO_PCI_CORE
+endif
 
 #
 # The conftest.sh script tests various aspects of the target kernel.
@@ -148,7 +158,11 @@ NV_CONFTEST_CMD := /bin/sh $(NV_CONFTEST_SCRIPT) \
 
 NV_CFLAGS_FROM_CONFTEST := $(shell $(NV_CONFTEST_CMD) build_cflags)
 
-NV_CONFTEST_CFLAGS = $(NV_CFLAGS_FROM_CONFTEST) $(EXTRA_CFLAGS) -fno-pie
+NV_CONFTEST_CFLAGS = $(NV_CFLAGS_FROM_CONFTEST) $(ccflags-y) -fno-pie
+NV_CONFTEST_CFLAGS += $(filter -std=%,$(KBUILD_CFLAGS))
+NV_CONFTEST_CFLAGS += $(call cc-disable-warning,pointer-sign)
+NV_CONFTEST_CFLAGS += $(call cc-option,-fshort-wchar,)
+NV_CONFTEST_CFLAGS += $(call cc-option,-Werror=incompatible-pointer-types,)
 
 NV_CONFTEST_COMPILE_TEST_HEADERS := $(obj)/conftest/macros.h
 NV_CONFTEST_COMPILE_TEST_HEADERS += $(obj)/conftest/functions.h
@@ -237,10 +251,12 @@ NV_HEADER_PRESENCE_TESTS = \
  drm/drm_device.h \
  drm/drm_mode_config.h \
  drm/drm_modeset_lock.h \
+ drm/drm_client_setup.h \
  dt-bindings/interconnect/tegra_icc_id.h \
  generated/autoconf.h \
  generated/compile.h \
  generated/utsrelease.h \
+ linux/aperture.h \
  linux/efi.h \
  linux/kconfig.h \
  linux/platform/tegra/mc_utils.h \
@@ -275,6 +291,7 @@ NV_HEADER_PRESENCE_TESTS = \
  asm/opal-api.h \
  sound/hdaudio.h \
  asm/pgtable_types.h \
+ asm/page.h \
  linux/stringhash.h \
  linux/dma-map-ops.h \
  rdma/peer_mem.h \
@@ -300,7 +317,10 @@ NV_HEADER_PRESENCE_TESTS = \
  linux/vfio_pci_core.h \
  linux/mdev.h \
  soc/tegra/bpmp-abi.h \
- soc/tegra/bpmp.h
+ soc/tegra/bpmp.h \
+ linux/cc_platform.h \
+ asm/cpufeature.h \
+ crypto/sig.h
 
 # Filename to store the define for the header in $(1); this is only consumed by
 # the rule below that concatenates all of these together.

kernel-open/Makefile

@@ -28,7 +28,7 @@ else
   else
     KERNEL_UNAME ?= $(shell uname -r)
     KERNEL_MODLIB := /lib/modules/$(KERNEL_UNAME)
-    KERNEL_SOURCES := $(shell test -d $(KERNEL_MODLIB)/source && echo $(KERNEL_MODLIB)/source || echo $(KERNEL_MODLIB)/build)
+    KERNEL_SOURCES := $(shell ((test -d $(KERNEL_MODLIB)/source && echo $(KERNEL_MODLIB)/source) || (test -d $(KERNEL_MODLIB)/build/source && echo $(KERNEL_MODLIB)/build/source)) || echo $(KERNEL_MODLIB)/build)
   endif
 
   KERNEL_OUTPUT := $(KERNEL_SOURCES)
@@ -42,12 +42,32 @@ else
   else
     KERNEL_UNAME ?= $(shell uname -r)
     KERNEL_MODLIB := /lib/modules/$(KERNEL_UNAME)
-    ifeq ($(KERNEL_SOURCES), $(KERNEL_MODLIB)/source)
+    # $(filter patter...,text) - Returns all whitespace-separated words in text that
+    # do match any of the pattern words, removing any words that do not match.
+    # Set the KERNEL_OUTPUT only if either $(KERNEL_MODLIB)/source or
+    # $(KERNEL_MODLIB)/build/source path matches the KERNEL_SOURCES.
+    ifneq ($(filter $(KERNEL_SOURCES),$(KERNEL_MODLIB)/source $(KERNEL_MODLIB)/build/source),)
       KERNEL_OUTPUT := $(KERNEL_MODLIB)/build
       KBUILD_PARAMS := KBUILD_OUTPUT=$(KERNEL_OUTPUT)
     endif
   endif
 
+  # If CC hasn't been set explicitly, check the value of CONFIG_CC_VERSION_TEXT.
+  # Look for the compiler specified there, and use it by default, if found.
+  ifeq ($(origin CC),default)
+    cc_version_text=$(firstword $(shell . $(KERNEL_OUTPUT)/.config; \
+                      echo "$$CONFIG_CC_VERSION_TEXT"))
+
+    ifneq ($(cc_version_text),)
+      ifeq ($(shell command -v $(cc_version_text)),)
+          $(warning WARNING: Unable to locate the compiler $(cc_version_text) \
+            from CONFIG_CC_VERSION_TEXT in the kernel configuration.)
+      else
+          CC=$(cc_version_text)
+      endif
+    endif
+  endif
+
   CC ?= cc
   LD ?= ld
   OBJDUMP ?= objdump
@@ -60,6 +80,16 @@ else
     )
   endif
 
+  KERNEL_ARCH = $(ARCH)
+
+  ifneq ($(filter $(ARCH),i386 x86_64),)
+    KERNEL_ARCH = x86
+  else
+    ifeq ($(filter $(ARCH),arm64 powerpc),)
+        $(error Unsupported architecture $(ARCH))
+    endif
+  endif
+
   NV_KERNEL_MODULES ?= $(wildcard nvidia nvidia-uvm nvidia-vgpu-vfio nvidia-modeset nvidia-drm nvidia-peermem)
   NV_KERNEL_MODULES := $(filter-out $(NV_EXCLUDE_KERNEL_MODULES), \
                                     $(NV_KERNEL_MODULES))
@@ -99,8 +129,9 @@ else
   # module symbols on which the Linux kernel's module resolution is dependent
   # and hence must be used whenever present.
 
-  LD_SCRIPT ?= $(KERNEL_SOURCES)/scripts/module-common.lds      \
-               $(KERNEL_SOURCES)/arch/$(ARCH)/kernel/module.lds \
+  LD_SCRIPT ?= $(KERNEL_SOURCES)/scripts/module-common.lds             \
+               $(KERNEL_SOURCES)/arch/$(KERNEL_ARCH)/kernel/module.lds \
+               $(KERNEL_OUTPUT)/arch/$(KERNEL_ARCH)/module.lds         \
                $(KERNEL_OUTPUT)/scripts/module.lds
   NV_MODULE_COMMON_SCRIPTS := $(foreach s, $(wildcard $(LD_SCRIPT)), -T $(s))
 

kernel-open/common/inc/nv-hypervisor.h

@@ -37,13 +37,11 @@ typedef enum _HYPERVISOR_TYPE
     OS_HYPERVISOR_UNKNOWN
 } HYPERVISOR_TYPE;
 
-#define CMD_VGPU_VFIO_WAKE_WAIT_QUEUE         0
-#define CMD_VGPU_VFIO_INJECT_INTERRUPT        1
-#define CMD_VGPU_VFIO_REGISTER_MDEV           2
-#define CMD_VGPU_VFIO_PRESENT                 3
-#define CMD_VFIO_PCI_CORE_PRESENT             4
+#define CMD_VFIO_WAKE_REMOVE_GPU              1
+#define CMD_VGPU_VFIO_PRESENT                 2
+#define CMD_VFIO_PCI_CORE_PRESENT             3
 
-#define MAX_VF_COUNT_PER_GPU 64
+#define MAX_VF_COUNT_PER_GPU                  64
 
 typedef enum _VGPU_TYPE_INFO
 {
@@ -54,17 +52,11 @@ typedef enum _VGPU_TYPE_INFO
 
 typedef struct
 {
-    void  *vgpuVfioRef;
-    void  *waitQueue;
     void  *nv;
-    NvU32 *vgpuTypeIds;
-    NvU8 **vgpuNames;
-    NvU32  numVgpuTypes;
-    NvU32  domain;
-    NvU8   bus;
-    NvU8   slot;
-    NvU8   function;
-    NvBool is_virtfn;
+    NvU32 domain;
+    NvU32 bus;
+    NvU32 device;
+    NvU32 return_status;
 } vgpu_vfio_info;
 
 typedef struct

kernel-open/common/inc/nv-linux.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2001-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2001-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -499,7 +499,9 @@ static inline void *nv_vmalloc(unsigned long size)
     void *ptr = __vmalloc(size, GFP_KERNEL);
 #endif
     if (ptr)
+    {
         NV_MEMDBG_ADD(ptr, size);
+    }
     return ptr;
 }
 
@@ -511,9 +513,15 @@ static inline void nv_vfree(void *ptr, NvU64 size)
 
 static inline void *nv_ioremap(NvU64 phys, NvU64 size)
 {
+#if IS_ENABLED(CONFIG_INTEL_TDX_GUEST) && defined(NV_IOREMAP_DRIVER_HARDENED_PRESENT)
+    void *ptr = ioremap_driver_hardened(phys, size);
+#else
     void *ptr = ioremap(phys, size);
+#endif
     if (ptr)
+    {
         NV_MEMDBG_ADD(ptr, size);
+    }
     return ptr;
 }
 
@@ -524,11 +532,11 @@ static inline void *nv_ioremap_nocache(NvU64 phys, NvU64 size)
 
 static inline void *nv_ioremap_cache(NvU64 phys, NvU64 size)
 {
-#if defined(NV_IOREMAP_CACHE_PRESENT)
-    void *ptr = ioremap_cache(phys, size);
-    if (ptr)
-        NV_MEMDBG_ADD(ptr, size);
-    return ptr;
+    void *ptr = NULL;
+#if IS_ENABLED(CONFIG_INTEL_TDX_GUEST) && defined(NV_IOREMAP_CACHE_SHARED_PRESENT)
+    ptr = ioremap_cache_shared(phys, size);
+#elif defined(NV_IOREMAP_CACHE_PRESENT)
+    ptr = ioremap_cache(phys, size);
 #elif defined(NVCPU_PPC64LE)
     //
     // ioremap_cache() has been only implemented correctly for ppc64le with
@@ -543,25 +551,34 @@ static inline void *nv_ioremap_cache(NvU64 phys, NvU64 size)
     // (commit 40f1ce7fb7e8, kernel 3.0+) and that covers all kernels we
     // support on power.
     //
-    void *ptr = ioremap_prot(phys, size, pgprot_val(PAGE_KERNEL));
-    if (ptr)
-        NV_MEMDBG_ADD(ptr, size);
-    return ptr;
+    ptr = ioremap_prot(phys, size, pgprot_val(PAGE_KERNEL));
 #else
     return nv_ioremap(phys, size);
 #endif
+
+    if (ptr)
+    {
+        NV_MEMDBG_ADD(ptr, size);
+    }
+    return ptr;
 }
 
 static inline void *nv_ioremap_wc(NvU64 phys, NvU64 size)
 {
-#if defined(NV_IOREMAP_WC_PRESENT)
-    void *ptr = ioremap_wc(phys, size);
-    if (ptr)
-        NV_MEMDBG_ADD(ptr, size);
-    return ptr;
+    void *ptr = NULL;
+#if IS_ENABLED(CONFIG_INTEL_TDX_GUEST) && defined(NV_IOREMAP_DRIVER_HARDENED_WC_PRESENT)
+    ptr = ioremap_driver_hardened_wc(phys, size);
+#elif defined(NV_IOREMAP_WC_PRESENT)
+    ptr = ioremap_wc(phys, size);
 #else
     return nv_ioremap_nocache(phys, size);
 #endif
+
+    if (ptr)
+    {
+        NV_MEMDBG_ADD(ptr, size);
+    }
+    return ptr;
 }
 
 static inline void nv_iounmap(void *ptr, NvU64 size)
@@ -634,37 +651,24 @@ static NvBool nv_numa_node_has_memory(int node_id)
         free_pages(ptr, order);                      \
     }
 
-extern NvU64 nv_shared_gpa_boundary;
+static inline pgprot_t nv_sme_clr(pgprot_t prot)
+{
+#if defined(__sme_clr)
+    return __pgprot(__sme_clr(pgprot_val(prot)));
+#else
+    return prot;
+#endif // __sme_clr
+}
 
 static inline pgprot_t nv_adjust_pgprot(pgprot_t vm_prot, NvU32 extra)
 {
     pgprot_t prot = __pgprot(pgprot_val(vm_prot) | extra);
-#if defined(CONFIG_AMD_MEM_ENCRYPT) && defined(NV_PGPROT_DECRYPTED_PRESENT)
-    /*
-     * When AMD memory encryption is enabled, device memory mappings with the
-     * C-bit set read as 0xFF, so ensure the bit is cleared for user mappings.
-     *
-     * If cc_mkdec() is present, then pgprot_decrypted() can't be used.
-     */
-#if defined(NV_CC_MKDEC_PRESENT)
-    if (nv_shared_gpa_boundary != 0)
-    {
-        /*
-         * By design, a VM using vTOM doesn't see the SEV setting and
-         * for AMD with vTOM, *set* means decrypted.
-         */
-        prot =  __pgprot(nv_shared_gpa_boundary | (pgprot_val(vm_prot)));
-    }
-    else
-    {
-        prot =  __pgprot(__sme_clr(pgprot_val(vm_prot)));
-    }
-#else
-    prot = pgprot_decrypted(prot);
-#endif
-#endif
 
-    return prot;
+#if defined(pgprot_decrypted)
+    return pgprot_decrypted(prot);
+#else
+    return nv_sme_clr(prot);
+#endif // pgprot_decrypted
 }
 
 #if defined(PAGE_KERNEL_NOENC)
@@ -702,7 +706,9 @@ static inline NvUPtr nv_vmap(struct page **pages, NvU32 page_count,
     /* All memory cached in PPC64LE; can't honor 'cached' input. */
     ptr = vmap(pages, page_count, VM_MAP, prot);
     if (ptr)
+    {
         NV_MEMDBG_ADD(ptr, page_count * PAGE_SIZE);
+    }
     return (NvUPtr)ptr;
 }
 
@@ -864,9 +870,9 @@ static inline dma_addr_t nv_phys_to_dma(struct device *dev, NvU64 pa)
 #define NV_PRINT_AT(nv_debug_level,at)                                           \
     {                                                                            \
         nv_printf(nv_debug_level,                                                \
-            "NVRM: VM: %s:%d: 0x%p, %d page(s), count = %d, flags = 0x%08x, "    \
+            "NVRM: VM: %s:%d: 0x%p, %d page(s), count = %lld, flags = 0x%08x, "  \
             "page_table = 0x%p\n",  __FUNCTION__, __LINE__, at,                  \
-            at->num_pages, NV_ATOMIC_READ(at->usage_count),                      \
+            at->num_pages, (long long)atomic64_read(&at->usage_count),           \
             at->flags, at->page_table);                                          \
     }
 
@@ -1190,7 +1196,7 @@ typedef struct nvidia_pte_s {
 typedef struct nv_alloc_s {
     struct nv_alloc_s *next;
     struct device     *dev;
-    atomic_t       usage_count;
+    atomic64_t       usage_count;
     struct {
         NvBool contig      : 1;
         NvBool guest       : 1;
@@ -1324,7 +1330,7 @@ nv_dma_maps_swiotlb(struct device *dev)
      * SEV memory encryption") forces SWIOTLB to be enabled when AMD SEV 
      * is active in all cases.
      */
-    if (os_sev_enabled)
+    if (os_cc_enabled)
         swiotlb_in_use = NV_TRUE;
 #endif
 
@@ -1487,7 +1493,8 @@ typedef struct
 typedef struct nv_linux_state_s {
     nv_state_t nv_state;
 
-    atomic_t usage_count;
+    atomic64_t usage_count;
+
     NvU32    suspend_count;
 
     struct device  *dev;
@@ -1605,6 +1612,10 @@ typedef struct nv_linux_state_s {
 
     struct nv_dma_device dma_dev;
     struct nv_dma_device niso_dma_dev;
+#if defined(NV_VGPU_KVM_BUILD)
+    wait_queue_head_t wait;
+    NvS32 return_status;
+#endif
 } nv_linux_state_t;
 
 extern nv_linux_state_t *nv_linux_devices;
@@ -1822,9 +1833,9 @@ static inline NvBool nv_alloc_release(nv_linux_file_private_t *nvlfp, nv_alloc_t
 {
     NV_PRINT_AT(NV_DBG_MEMINFO, at);
 
-    if (NV_ATOMIC_DEC_AND_TEST(at->usage_count))
+    if (atomic64_dec_and_test(&at->usage_count))
     {
-        NV_ATOMIC_INC(at->usage_count);
+        atomic64_inc(&at->usage_count);
 
         at->next = nvlfp->free_list;
         nvlfp->free_list = at;
@@ -1984,31 +1995,6 @@ static inline NvBool nv_platform_use_auto_online(nv_linux_state_t *nvl)
     return nvl->numa_info.use_auto_online;
 }
 
-typedef struct {
-    NvU64 base;
-    NvU64 size;
-    NvU32 nodeId;
-    int ret;
-} remove_numa_memory_info_t;
-
-static void offline_numa_memory_callback
-(
-    void *args
-)
-{
-#ifdef NV_OFFLINE_AND_REMOVE_MEMORY_PRESENT
-    remove_numa_memory_info_t *pNumaInfo = (remove_numa_memory_info_t *)args;
-#ifdef NV_REMOVE_MEMORY_HAS_NID_ARG
-    pNumaInfo->ret = offline_and_remove_memory(pNumaInfo->nodeId,
-                                               pNumaInfo->base,
-                                               pNumaInfo->size);
-#else
-    pNumaInfo->ret = offline_and_remove_memory(pNumaInfo->base,
-                                               pNumaInfo->size);
-#endif
-#endif
-}
-
 typedef enum
 {
     NV_NUMA_STATUS_DISABLED             = 0,
@@ -2069,4 +2055,7 @@ typedef enum
 #include <linux/clk-provider.h>
 #endif
 
+#define NV_EXPORT_SYMBOL(symbol)        EXPORT_SYMBOL_GPL(symbol)
+#define NV_CHECK_EXPORT_SYMBOL(symbol)  NV_IS_EXPORT_SYMBOL_PRESENT_##symbol
+
 #endif  /* _NV_LINUX_H_ */

kernel-open/common/inc/nv-lock.h

@@ -35,17 +35,6 @@
 #include <linux/sched/signal.h>     /* signal_pending for kernels >= 4.11 */
 #endif
 
-#if defined(CONFIG_PREEMPT_RT) || defined(CONFIG_PREEMPT_RT_FULL)
-typedef raw_spinlock_t            nv_spinlock_t;
-#define NV_SPIN_LOCK_INIT(lock)   raw_spin_lock_init(lock)
-#define NV_SPIN_LOCK_IRQ(lock)    raw_spin_lock_irq(lock)
-#define NV_SPIN_UNLOCK_IRQ(lock)  raw_spin_unlock_irq(lock)
-#define NV_SPIN_LOCK_IRQSAVE(lock,flags) raw_spin_lock_irqsave(lock,flags)
-#define NV_SPIN_UNLOCK_IRQRESTORE(lock,flags) raw_spin_unlock_irqrestore(lock,flags)
-#define NV_SPIN_LOCK(lock)        raw_spin_lock(lock)
-#define NV_SPIN_UNLOCK(lock)      raw_spin_unlock(lock)
-#define NV_SPIN_UNLOCK_WAIT(lock) raw_spin_unlock_wait(lock)
-#else
 typedef spinlock_t                nv_spinlock_t;
 #define NV_SPIN_LOCK_INIT(lock)   spin_lock_init(lock)
 #define NV_SPIN_LOCK_IRQ(lock)    spin_lock_irq(lock)
@@ -55,7 +44,6 @@ typedef spinlock_t                nv_spinlock_t;
 #define NV_SPIN_LOCK(lock)        spin_lock(lock)
 #define NV_SPIN_UNLOCK(lock)      spin_unlock(lock)
 #define NV_SPIN_UNLOCK_WAIT(lock) spin_unlock_wait(lock)
-#endif
 
 #define NV_INIT_MUTEX(mutex) sema_init(mutex, 1)
 

kernel-open/common/inc/nv-mm.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2016-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2016-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -292,9 +292,21 @@ static inline struct rw_semaphore *nv_mmap_get_lock(struct mm_struct *mm)
 #endif
 }
 
+#define NV_CAN_CALL_VMA_START_WRITE 1
+
+#if !NV_CAN_CALL_VMA_START_WRITE
+/*
+ * Commit 45ad9f5290dc updated vma_start_write() to call __vma_start_write().
+ */
+void nv_vma_start_write(struct vm_area_struct *);
+#endif
+
 static inline void nv_vm_flags_set(struct vm_area_struct *vma, vm_flags_t flags)
 {
-#if defined(NV_VM_AREA_STRUCT_HAS_CONST_VM_FLAGS)
+#if !NV_CAN_CALL_VMA_START_WRITE
+    nv_vma_start_write(vma);
+    ACCESS_PRIVATE(vma, __vm_flags) |= flags;
+#elif defined(NV_VM_AREA_STRUCT_HAS_CONST_VM_FLAGS)
     vm_flags_set(vma, flags);
 #else
     vma->vm_flags |= flags;
@@ -303,7 +315,10 @@ static inline void nv_vm_flags_set(struct vm_area_struct *vma, vm_flags_t flags)
 
 static inline void nv_vm_flags_clear(struct vm_area_struct *vma, vm_flags_t flags)
 {
-#if defined(NV_VM_AREA_STRUCT_HAS_CONST_VM_FLAGS)
+#if !NV_CAN_CALL_VMA_START_WRITE
+    nv_vma_start_write(vma);
+    ACCESS_PRIVATE(vma, __vm_flags) &= ~flags;
+#elif defined(NV_VM_AREA_STRUCT_HAS_CONST_VM_FLAGS)
     vm_flags_clear(vma, flags);
 #else
     vma->vm_flags &= ~flags;

kernel-open/common/inc/nv-timer.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2017 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2017-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -63,4 +63,13 @@ static inline void nv_timer_setup(struct nv_timer *nv_timer,
 #endif
 }
 
+static inline void nv_timer_delete_sync(struct timer_list *timer)
+{
+#if !defined(NV_BSD) && NV_IS_EXPORT_SYMBOL_PRESENT_timer_delete_sync
+    timer_delete_sync(timer);
+#else
+    del_timer_sync(timer);
+#endif
+}
+
 #endif // __NV_TIMER_H__

kernel-open/common/inc/nv.h

@@ -615,6 +615,14 @@ typedef enum
 #define NV_IS_DEVICE_IN_SURPRISE_REMOVAL(nv)    \
         (((nv)->flags & NV_FLAG_IN_SURPRISE_REMOVAL) != 0)
 
+/*
+ * For console setup by EFI GOP, the base address is BAR1.
+ * For console setup by VBIOS, the base address is BAR2 + 16MB.
+ */
+#define NV_IS_CONSOLE_MAPPED(nv, addr)  \
+        (((addr) == (nv)->bars[NV_GPU_BAR_INDEX_FB].cpu_address) || \
+         ((addr) == ((nv)->bars[NV_GPU_BAR_INDEX_IMEM].cpu_address + 0x1000000)))
+
 #define NV_SOC_IS_ISO_IOMMU_PRESENT(nv)     \
         ((nv)->iso_iommu_present)
 
@@ -874,6 +882,8 @@ NvBool    NV_API_CALL nv_match_gpu_os_info(nv_state_t *, void *);
 NvU32     NV_API_CALL nv_get_os_type(void);
 
 void      NV_API_CALL nv_get_updated_emu_seg(NvU32 *start, NvU32 *end);
+void      NV_API_CALL nv_get_screen_info(nv_state_t *, NvU64 *, NvU16 *, NvU16 *, NvU16 *, NvU16 *, NvU64 *);
+
 struct dma_buf;
 typedef struct nv_dma_buf nv_dma_buf_t;
 struct drm_gem_object;
@@ -924,6 +934,7 @@ NV_STATUS  NV_API_CALL  rm_ioctl                 (nvidia_stack_t *, nv_state_t *
 NvBool     NV_API_CALL  rm_isr                   (nvidia_stack_t *, nv_state_t *, NvU32 *);
 void       NV_API_CALL  rm_isr_bh                (nvidia_stack_t *, nv_state_t *);
 void       NV_API_CALL  rm_isr_bh_unlocked       (nvidia_stack_t *, nv_state_t *);
+NvBool     NV_API_CALL  rm_is_msix_allowed       (nvidia_stack_t *, nv_state_t *);
 NV_STATUS  NV_API_CALL  rm_power_management      (nvidia_stack_t *, nv_state_t *, nv_pm_action_t);
 NV_STATUS  NV_API_CALL  rm_stop_user_channels    (nvidia_stack_t *, nv_state_t *);
 NV_STATUS  NV_API_CALL  rm_restart_user_channels (nvidia_stack_t *, nv_state_t *);
@@ -1023,12 +1034,11 @@ NV_STATUS  NV_API_CALL  nv_vgpu_create_request(nvidia_stack_t *, nv_state_t *, c
 NV_STATUS  NV_API_CALL  nv_vgpu_delete(nvidia_stack_t *, const NvU8 *, NvU16);
 NV_STATUS  NV_API_CALL  nv_vgpu_get_type_ids(nvidia_stack_t *, nv_state_t *, NvU32 *, NvU32 *, NvBool, NvU8, NvBool);
 NV_STATUS  NV_API_CALL  nv_vgpu_get_type_info(nvidia_stack_t *, nv_state_t *, NvU32, char *, int, NvU8);
-NV_STATUS  NV_API_CALL  nv_vgpu_get_bar_info(nvidia_stack_t *, nv_state_t *, const NvU8 *, NvU64 *, NvU32, void *);
-NV_STATUS  NV_API_CALL  nv_vgpu_start(nvidia_stack_t *, const NvU8 *, void *, NvS32 *, NvU8 *, NvU32);
-NV_STATUS  NV_API_CALL  nv_vgpu_get_sparse_mmap(nvidia_stack_t *, nv_state_t *, const NvU8 *, NvU64 **, NvU64 **, NvU32 *);
+NV_STATUS  NV_API_CALL  nv_vgpu_get_bar_info(nvidia_stack_t *, nv_state_t *, const NvU8 *, NvU64 *,
+                                             NvU64 *, NvU64 *, NvU32 *, NvU8 *);
 NV_STATUS  NV_API_CALL  nv_vgpu_process_vf_info(nvidia_stack_t *, nv_state_t *, NvU8, NvU32, NvU8, NvU8, NvU8, NvBool, void *);
-NV_STATUS  NV_API_CALL  nv_vgpu_update_request(nvidia_stack_t *, const NvU8 *, NvU32, NvU64 *, NvU64 *, const char *);
-NV_STATUS  NV_API_CALL  nv_gpu_bind_event(nvidia_stack_t *);
+NV_STATUS  NV_API_CALL  nv_gpu_bind_event(nvidia_stack_t *, NvU32, NvBool *);
+NV_STATUS  NV_API_CALL  nv_gpu_unbind_event(nvidia_stack_t *, NvU32, NvBool *);
 
 NV_STATUS NV_API_CALL nv_get_usermap_access_params(nv_state_t*, nv_usermap_access_params_t*);
 nv_soc_irq_type_t NV_API_CALL nv_get_current_irq_type(nv_state_t*);

kernel-open/common/inc/nv_uvm_types.h

@@ -321,10 +321,6 @@ typedef struct UvmGpuChannelAllocParams_tag
     // The next two fields store UVM_BUFFER_LOCATION values
     NvU32 gpFifoLoc;
     NvU32 gpPutLoc;
-
-    // Allocate the channel as secure. This flag should only be set when
-    // Confidential Compute is enabled.
-    NvBool secure;
 } UvmGpuChannelAllocParams;
 
 typedef struct UvmGpuPagingChannelAllocParams_tag
@@ -368,9 +364,6 @@ typedef struct
     // True if the CE can be used for P2P transactions
     NvBool p2p:1;
 
-    // True if the CE supports encryption
-    NvBool secure:1;
-
     // Mask of physical CEs assigned to this LCE
     //
     // The value returned by RM for this field may change when a GPU is
@@ -573,8 +566,11 @@ typedef struct UvmPlatformInfo_tag
     // Out: ATS (Address Translation Services) is supported
     NvBool atsSupported;
 
-    // Out: AMD SEV (Secure Encrypted Virtualization) is enabled
-    NvBool sevEnabled;
+    // Out: True if HW trusted execution, such as AMD's SEV-SNP or Intel's TDX,
+    // is enabled in the VM, indicating that Confidential Computing must be
+    // also enabled in the GPU(s); these two security features are either both
+    // enabled, or both disabled.
+    NvBool confComputingEnabled;
 } UvmPlatformInfo;
 
 typedef struct UvmGpuClientInfo_tag

kernel-open/common/inc/nvmisc.h

@@ -694,6 +694,42 @@ nvPrevPow2_U64(const NvU64 x )
     }                                                       \
 }
 
+//
+// Bug 4851259: Newly added functions must be hidden from certain HS-signed
+// ucode compilers to avoid signature mismatch.
+//
+#ifndef NVDEC_1_0
+/*!
+ * Returns the position of nth set bit in the given mask.
+ *
+ * Returns -1 if mask has fewer than n bits set.
+ *
+ * n is 0 indexed and has valid values 0..31 inclusive, so "zeroth" set bit is
+ * the first set LSB.
+ *
+ * Example, if mask = 0x000000F0u and n = 1, the return value will be 5.
+ * Example, if mask = 0x000000F0u and n = 4, the return value will be -1.
+ */
+static NV_FORCEINLINE NvS32
+nvGetNthSetBitIndex32(NvU32 mask, NvU32 n)
+{
+    NvU32 seenSetBitsCount = 0;
+    NvS32 index;
+    FOR_EACH_INDEX_IN_MASK(32, index, mask)
+    {
+        if (seenSetBitsCount == n)
+        {
+            return index;
+        }
+        ++seenSetBitsCount;
+    }
+    FOR_EACH_INDEX_IN_MASK_END;
+
+    return -1;
+}
+
+#endif // NVDEC_1_0
+
 //
 // Size to use when declaring variable-sized arrays
 //

kernel-open/common/inc/os-interface.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 1999-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 1999-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -162,10 +162,9 @@ NvBool      NV_API_CALL  os_is_vgx_hyper             (void);
 NV_STATUS   NV_API_CALL  os_inject_vgx_msi           (NvU16, NvU64, NvU32);
 NvBool      NV_API_CALL  os_is_grid_supported        (void);
 NvU32       NV_API_CALL  os_get_grid_csp_support     (void);
-void        NV_API_CALL  os_get_screen_info          (NvU64 *, NvU16 *, NvU16 *, NvU16 *, NvU16 *, NvU64, NvU64);
 void        NV_API_CALL  os_bug_check                (NvU32, const char *);
 NV_STATUS   NV_API_CALL  os_lock_user_pages          (void *, NvU64, void **, NvU32);
-NV_STATUS   NV_API_CALL  os_lookup_user_io_memory    (void *, NvU64, NvU64 **, void**);
+NV_STATUS   NV_API_CALL  os_lookup_user_io_memory    (void *, NvU64, NvU64 **);
 NV_STATUS   NV_API_CALL  os_unlock_user_pages        (NvU64, void *);
 NV_STATUS   NV_API_CALL  os_match_mmap_offset        (void *, NvU64, NvU64 *);
 NV_STATUS   NV_API_CALL  os_get_euid                 (NvU32 *);
@@ -207,15 +206,19 @@ enum os_pci_req_atomics_type {
     OS_INTF_PCIE_REQ_ATOMICS_128BIT
 };
 NV_STATUS   NV_API_CALL  os_enable_pci_req_atomics   (void *, enum os_pci_req_atomics_type);
+NV_STATUS   NV_API_CALL  os_get_numa_node_memory_usage (NvS32, NvU64 *, NvU64 *);
 NV_STATUS   NV_API_CALL  os_numa_add_gpu_memory      (void *, NvU64, NvU64, NvU32 *);
 NV_STATUS   NV_API_CALL  os_numa_remove_gpu_memory   (void *, NvU64, NvU64, NvU32); 
 NV_STATUS   NV_API_CALL  os_offline_page_at_address(NvU64 address);
+void*       NV_API_CALL  os_get_pid_info(void);
+void        NV_API_CALL  os_put_pid_info(void *pid_info);
+NV_STATUS   NV_API_CALL  os_find_ns_pid(void *pid_info, NvU32 *ns_pid);
 
 extern NvU32 os_page_size;
 extern NvU64 os_page_mask;
 extern NvU8  os_page_shift;
-extern NvU32 os_sev_status;
-extern NvBool os_sev_enabled;
+extern NvBool os_cc_enabled;
+extern NvBool os_cc_tdx_enabled;
 extern NvBool os_dma_buf_enabled;
 
 /*

kernel-open/conftest.sh

@@ -316,7 +316,7 @@ export_symbol_present_conftest() {
     SYMBOL="$1"
     TAB='	'
 
-    if grep -e "${TAB}${SYMBOL}${TAB}.*${TAB}EXPORT_SYMBOL.*\$" \
+    if grep -e "${TAB}${SYMBOL}${TAB}.*${TAB}EXPORT_SYMBOL\(_GPL\)\?\s*\$" \
                "$OUTPUT/Module.symvers" >/dev/null 2>&1; then
         echo "#define NV_IS_EXPORT_SYMBOL_PRESENT_$SYMBOL 1" |
             append_conftest "symbols"
@@ -337,7 +337,7 @@ export_symbol_gpl_conftest() {
     SYMBOL="$1"
     TAB='	'
 
-    if grep -e "${TAB}${SYMBOL}${TAB}.*${TAB}EXPORT_\(UNUSED_\)*SYMBOL_GPL\$" \
+    if grep -e "${TAB}${SYMBOL}${TAB}.*${TAB}EXPORT_\(UNUSED_\)*SYMBOL_GPL\s*\$" \
                "$OUTPUT/Module.symvers" >/dev/null 2>&1; then
         echo "#define NV_IS_EXPORT_SYMBOL_GPL_$SYMBOL 1" |
             append_conftest "symbols"
@@ -445,6 +445,9 @@ compile_test() {
             #if defined(NV_ASM_PGTABLE_TYPES_H_PRESENT)
             #include <asm/pgtable_types.h>
             #endif
+            #if defined(NV_ASM_PAGE_H_PRESENT)
+            #include <asm/page.h>
+            #endif
             #include <asm/set_memory.h>
             #else
             #include <asm/cacheflush.h>
@@ -467,6 +470,9 @@ compile_test() {
             #if defined(NV_ASM_PGTABLE_TYPES_H_PRESENT)
             #include <asm/pgtable_types.h>
             #endif
+            #if defined(NV_ASM_PAGE_H_PRESENT)
+            #include <asm/page.h>
+            #endif
             #include <asm/set_memory.h>
             #else
             #include <asm/cacheflush.h>
@@ -524,6 +530,9 @@ compile_test() {
             #if defined(NV_ASM_PGTABLE_TYPES_H_PRESENT)
             #include <asm/pgtable_types.h>
             #endif
+            #if defined(NV_ASM_PAGE_H_PRESENT)
+            #include <asm/page.h>
+            #endif
             #include <asm/set_memory.h>
             #else
             #include <asm/cacheflush.h>
@@ -551,6 +560,9 @@ compile_test() {
             #if defined(NV_ASM_PGTABLE_TYPES_H_PRESENT)
             #include <asm/pgtable_types.h>
             #endif
+            #if defined(NV_ASM_PAGE_H_PRESENT)
+            #include <asm/page.h>
+            #endif
             #include <asm/set_memory.h>
             #else
             #include <asm/cacheflush.h>
@@ -695,6 +707,50 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_IOREMAP_WC_PRESENT" "" "functions"
         ;;
 
+        ioremap_driver_hardened)
+            #
+            # Determine if the ioremap_driver_hardened() function is present.
+            # It does not exist on all architectures.
+            # TODO: Update the commit ID once the API is upstreamed.
+            #
+            CODE="
+            #include <asm/io.h>
+            void conftest_ioremap_driver_hardened(void) {
+                ioremap_driver_hardened();
+            }"
+
+            compile_check_conftest "$CODE" "NV_IOREMAP_DRIVER_HARDENED_PRESENT" "" "functions"
+        ;;
+
+        ioremap_driver_hardened_wc)
+            #
+            # Determine if the ioremap_driver_hardened_wc() function is present.
+            # It does not exist on all architectures.
+            # TODO: Update the commit ID once the API is upstreamed.
+            #
+            CODE="
+            #include <asm/io.h>
+            void conftest_ioremap_driver_hardened_wc(void) {
+                ioremap_driver_hardened_wc();
+            }"
+
+            compile_check_conftest "$CODE" "NV_IOREMAP_DRIVER_HARDENED_WC_PRESENT" "" "functions"
+        ;;
+
+        ioremap_cache_shared)
+            #
+            # Determine if the ioremap_cache_shared() function is present.
+            # It does not exist on all architectures.
+            # TODO: Update the commit ID once the API is upstreamed.
+            #
+            CODE="
+            #include <asm/io.h>
+            void conftest_ioremap_cache_shared(void) {
+                ioremap_cache_shared();
+            }"
+
+            compile_check_conftest "$CODE" "NV_IOREMAP_CACHE_SHARED_PRESENT" "" "functions"
+        ;;
         dom0_kernel_present)
             # Add config parameter if running on DOM0.
             if [ -n "$VGX_BUILD" ]; then
@@ -1152,6 +1208,23 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_VFIO_DEVICE_OPS_HAS_BIND_IOMMUFD" "" "types"
         ;;
 
+        vfio_device_ops_has_detach_ioas)
+            #
+            # Determine if 'vfio_device_ops' struct has 'detach_ioas' field.
+            #
+            # Added by commit 9048c7341c4df9cae04c154a8b0f556dbe913358 ("vfio-iommufd: Add detach_ioas
+            # support for physical VFIO devices
+            #
+            CODE="
+            #include <linux/pci.h>
+            #include <linux/vfio.h>
+            int conftest_vfio_device_ops_has_detach_ioas(void) {
+                return offsetof(struct vfio_device_ops, detach_ioas);
+            }"
+
+            compile_check_conftest "$CODE" "NV_VFIO_DEVICE_OPS_HAS_DETACH_IOAS" "" "types"
+        ;;
+
         pci_irq_vector_helpers)
             #
             # Determine if pci_alloc_irq_vectors(), pci_free_irq_vectors()
@@ -1298,6 +1371,42 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_VFIO_REGISTER_EMULATED_IOMMU_DEV_PRESENT" "" "functions"
         ;;
 
+        bus_type_has_iommu_ops)
+            #
+            # Determine if 'bus_type' structure has a 'iommu_ops' field.
+            #
+            # This field was removed by commit 17de3f5fdd35 (iommu: Retire bus ops)
+            # in v6.8
+            #
+            CODE="
+            #include <linux/device.h>
+
+            int conftest_bus_type_has_iommu_ops(void) {
+                return offsetof(struct bus_type, iommu_ops);
+            }"
+
+            compile_check_conftest "$CODE" "NV_BUS_TYPE_HAS_IOMMU_OPS" "" "types"
+        ;;
+
+        eventfd_signal_has_counter_arg)
+            #
+            # Determine if eventfd_signal() function has an additional 'counter' argument.
+            #
+            # This argument was removed by commit 3652117f8548 (eventfd: simplify
+            # eventfd_signal()) in v6.8
+            #
+            CODE="
+            #include <linux/eventfd.h>
+
+            void conftest_eventfd_signal_has_counter_arg(void) {
+                struct eventfd_ctx *ctx;
+
+                eventfd_signal(ctx, 1);
+            }"
+
+            compile_check_conftest "$CODE" "NV_EVENTFD_SIGNAL_HAS_COUNTER_ARG" "" "types"
+        ;;
+
         drm_available)
             # Determine if the DRM subsystem is usable
             CODE="
@@ -2366,6 +2475,22 @@ compile_test() {
             fi
         ;;
 
+        file_operations_fop_unsigned_offset_present)
+            #
+            # Determine if the FOP_UNSIGNED_OFFSET define is present.
+            #
+            # Added by commit 641bb4394f40 ("fs: move FMODE_UNSIGNED_OFFSET to
+            # fop_flags") in v6.12.
+            #
+            CODE="
+            #include <linux/fs.h>
+            int conftest_file_operations_fop_unsigned_offset_present(void) {
+                return FOP_UNSIGNED_OFFSET;
+            }"
+
+            compile_check_conftest "$CODE" "NV_FILE_OPERATIONS_FOP_UNSIGNED_OFFSET_PRESENT" "" "types"
+        ;;
+
         mm_context_t)
             #
             # Determine if the 'mm_context_t' data type is present
@@ -2976,6 +3101,22 @@ compile_test() {
 
         ;;
 
+        foll_longterm_present)
+            #
+            # Determine if FOLL_LONGTERM enum is present or not
+            #
+            # Added by commit 932f4a630a69 ("mm/gup: replace
+            # get_user_pages_longterm() with FOLL_LONGTERM") in
+            # v5.2
+            #
+            CODE="
+            #include <linux/mm.h>
+            int foll_longterm = FOLL_LONGTERM;
+            "
+
+            compile_check_conftest "$CODE" "NV_FOLL_LONGTERM_PRESENT" "" "types"
+        ;;
+
         vfio_pin_pages_has_vfio_device_arg)
             #
             # Determine if vfio_pin_pages() kABI accepts "struct vfio_device *"
@@ -3872,33 +4013,6 @@ compile_test() {
             fi
         ;;
 
-        dma_buf_has_dynamic_attachment)
-            #
-            # Determine if the function dma_buf_attachment_is_dynamic()
-            # is present.
-            #
-            # Added by commit: 15fd552d186c
-            # ("dma-buf: change DMA-buf locking convention v3") in v5.5 (2018-07-03)
-            #
-            echo "$CONFTEST_PREAMBLE
-            #include <linux/dma-buf.h>
-            bool conftest_dma_buf_attachment_is_dynamic(void) {
-                return dma_buf_attachment_is_dynamic(NULL);
-            }" > conftest$$.c
-
-            $CC $CFLAGS -c conftest$$.c > /dev/null 2>&1
-            rm -f conftest$$.c
-
-            if [ -f conftest$$.o ]; then
-                echo "#define NV_DMA_BUF_HAS_DYNAMIC_ATTACHMENT" | append_conftest "functions"
-                rm -f conftest$$.o
-                return
-            else
-                echo "#undef NV_DMA_BUF_HAS_DYNAMIC_ATTACHMENT" | append_conftest "functions"
-                return
-            fi
-        ;;
-
         dma_buf_attachment_has_peer2peer)
             #
             # Determine if peer2peer is present in struct dma_buf_attachment.
@@ -3927,6 +4041,43 @@ compile_test() {
             fi
         ;;
 
+        drm_fb_create_takes_format_info)
+            #
+            # Determine if a `struct drm_format_info *` is passed into
+            # the .fb_create callback. If so, it will have 4 arguments.
+            # This parameter was added in commit 81112eaac559 ("drm:
+            # Pass the format info to .fb_create") in linux-next
+            # (2025-07-16)
+            CODE="
+            #include <drm/drm_mode_config.h>
+            #include <drm/drm_framebuffer.h>
+
+            static const struct drm_mode_config_funcs funcs;
+            void conftest_drm_fb_create_takes_format_info(void) {
+                funcs.fb_create(NULL, NULL, NULL, NULL); 
+            }"
+
+            compile_check_conftest "$CODE" "NV_DRM_FB_CREATE_TAKES_FORMAT_INFO" "" "types"
+        ;;
+
+        drm_fill_fb_struct_takes_format_info)
+            #
+            # Determine if a `struct drm_format_info *` is passed into
+            # drm_helper_mode_fill_fb_struct(). If so, it will have 4 arguments.
+            # This parameter was added in commit a34cc7bf1034 ("drm:
+            # Allow the caller to pass in the format info to
+            # drm_helper_mode_fill_fb_struct()") in linux-next
+            # (2025-07-16)
+            CODE="
+            #include <drm/drm_modeset_helper.h>
+
+            void conftest_drm_fill_fb_struct_takes_format_info(void) {
+                drm_helper_mode_fill_fb_struct(NULL, NULL, NULL, NULL);
+            }"
+
+            compile_check_conftest "$CODE" "NV_DRM_FILL_FB_STRUCT_TAKES_FORMAT_INFO" "" "types"
+        ;;
+
         drm_connector_funcs_have_mode_in_name)
             #
             # Determine if _mode_ is present in connector function names.  We
@@ -4412,6 +4563,24 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_MMU_NOTIFIER_OPS_HAS_INVALIDATE_RANGE" "" "types"
         ;;
 
+        mmu_notifier_ops_arch_invalidate_secondary_tlbs)
+            #
+            # Determine if the mmu_notifier_ops struct has the
+            # 'arch_invalidate_secondary_tlbs' member.
+            #
+            # struct mmu_notifier_ops.invalidate_range was renamed to
+            # arch_invalidate_secondary_tlbs by commit 1af5a8109904
+            # ("mmu_notifiers: rename invalidate_range notifier") due to be
+            # added in v6.6
+           CODE="
+            #include <linux/mmu_notifier.h>
+            int conftest_mmu_notifier_ops_arch_invalidate_secondary_tlbs(void) {
+                return offsetof(struct mmu_notifier_ops, arch_invalidate_secondary_tlbs);
+            }"
+
+            compile_check_conftest "$CODE" "NV_MMU_NOTIFIER_OPS_HAS_ARCH_INVALIDATE_SECONDARY_TLBS" "" "types"
+        ;;
+
         drm_format_num_planes)
             #
             # Determine if drm_format_num_planes() function is present.
@@ -4888,40 +5057,22 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_PCI_CHANNEL_STATE_PRESENT" "" "types"
         ;;
 
-        pgprot_decrypted)
+        cc_platform_has)
             #
-            # Determine if the macro 'pgprot_decrypted()' is present.
+            # Determine if 'cc_platform_has()' is present.
             #
-            # Added by commit 21729f81ce8a ("x86/mm: Provide general kernel
-            # support for memory encryption") in v4.14 (2017-07-18)
+            # Added by commit aa5a461171f9 ("x86/sev: Add an x86 version of
+            # cc_platform_has()") in v5.15.3 (2021-10-04)
             CODE="
-            #include <asm/pgtable.h>
-
-            void conftest_pgprot_decrypted(void)
-                if(pgprot_decrypted()) {}
-            }"
-
-            compile_check_conftest "$CODE" "NV_PGPROT_DECRYPTED_PRESENT" "" "functions"
-
-        ;;
-
-        cc_mkdec)
-            #
-            # Determine if cc_mkdec() is present.
-            #
-            # cc_mkdec() by commit b577f542f93c ("x86/coco: Add API to handle
-            # encryption mask) in v5.18-rc1 (2022-02-22).
-            #
-            CODE="
-            #if defined(NV_ASM_COCO_H_PRESENT)
-            #include <asm/coco.h>
+            #if defined(NV_LINUX_CC_PLATFORM_H_PRESENT)
+            #include <linux/cc_platform.h>
             #endif
 
-            void conftest_cc_mkdec(void) {
-                cc_mkdec();
+            void conftest_cc_platfrom_has(void) {
+                cc_platform_has();
             }"
 
-            compile_check_conftest "$CODE" "NV_CC_MKDEC_PRESENT" "" "functions"
+            compile_check_conftest "$CODE" "NV_CC_PLATFORM_PRESENT" "" "functions"
         ;;
 
         drm_prime_pages_to_sg_has_drm_device_arg)
@@ -5025,11 +5176,15 @@ compile_test() {
             # vmap ops and convert GEM backends") update
             # drm_gem_object_funcs::vmap to take 'map' argument.
             #
+            # Note that the 'map' argument type is changed from 'struct dma_buf_map'
+            # to 'struct iosys_map' by commit 7938f4218168 ("dma-buf-map: Rename
+            # to iosys-map) in v5.18.
+            #
             CODE="
             #include <drm/drm_gem.h>
             int conftest_drm_gem_object_vmap_has_map_arg(
-                    struct drm_gem_object *obj, struct dma_buf_map *map) {
-                return obj->funcs->vmap(obj, map);
+                    struct drm_gem_object *obj) {
+                return obj->funcs->vmap(obj, NULL);
             }"
 
             compile_check_conftest "$CODE" "NV_DRM_GEM_OBJECT_VMAP_HAS_MAP_ARG" "" "types"
@@ -5069,20 +5224,60 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_PCI_CLASS_MULTIMEDIA_HD_AUDIO_PRESENT" "" "generic"
         ;;
 
-        unsafe_follow_pfn)
+        follow_pfn)
             #
-            # Determine if unsafe_follow_pfn() is present.
+            # Determine if follow_pfn() is present.
             #
-            # unsafe_follow_pfn() was added by commit 69bacee7f9ad
-            # ("mm: Add unsafe_follow_pfn") in v5.13-rc1.
+            # follow_pfn() was added by commit 3b6748e2dd69
+            # ("mm: introduce follow_pfn()") in v2.6.31-rc1, and removed
+            # by commit 233eb0bf3b94 ("mm: remove follow_pfn")
+            # from linux-next 233eb0bf3b94.
             #
             CODE="
             #include <linux/mm.h>
-            void conftest_unsafe_follow_pfn(void) {
-                unsafe_follow_pfn();
+            void conftest_follow_pfn(void) {
+                follow_pfn();
             }"
 
-            compile_check_conftest "$CODE" "NV_UNSAFE_FOLLOW_PFN_PRESENT" "" "functions"
+            compile_check_conftest "$CODE" "NV_FOLLOW_PFN_PRESENT" "" "functions"
+        ;;
+
+        follow_pte_arg_vma)
+            #
+            # Determine if the first argument of follow_pte is
+            # mm_struct or vm_area_struct.
+            #
+            # The first argument was changed from mm_struct to vm_area_struct by
+            # commit 29ae7d96d166 ("mm: pass VMA instead of MM to follow_pte()")
+            #
+            CODE="
+            #include <linux/mm.h>
+
+            typeof(follow_pte) conftest_follow_pte_has_vma_arg;
+            int conftest_follow_pte_has_vma_arg(struct vm_area_struct *vma,
+                                                unsigned long address,
+                                                pte_t **ptep,
+                                                spinlock_t **ptl) {
+                return 0;
+            }"
+
+            compile_check_conftest "$CODE" "NV_FOLLOW_PTE_ARG1_VMA" "" "types"
+        ;;
+
+        ptep_get)
+            #
+            # Determine if ptep_get() is present.
+            #
+            # ptep_get() was added by commit 481e980a7c19
+            # ("mm: Allow arches to provide ptep_get()")
+            #
+            CODE="
+            #include <linux/mm.h>
+            void conftest_ptep_get(void) {
+                ptep_get();
+            }"
+
+            compile_check_conftest "$CODE" "NV_PTEP_GET_PRESENT" "" "functions"
         ;;
 
         drm_plane_atomic_check_has_atomic_state_arg)
@@ -5157,10 +5352,16 @@ compile_test() {
             # Added by commit 7b7b27214bba ("mm/memory_hotplug: introduce
             # add_memory_driver_managed()") in v5.8-rc1 (2020-06-05)
             #
+            # Before commit 3a0aaefe4134 ("mm/memory_hotplug: guard more
+            # declarations by CONFIG_MEMORY_HOTPLUG") in v5.10, the
+            # add_memory_driver_managed() was not guarded.
+            #
             CODE="
             #include <linux/memory_hotplug.h>
             void conftest_add_memory_driver_managed() {
+            #if defined(CONFIG_MEMORY_HOTPLUG)
                 add_memory_driver_managed();
+            #endif
             }"
 
             compile_check_conftest "$CODE" "NV_ADD_MEMORY_DRIVER_MANAGED_PRESENT" "" "functions"
@@ -5598,23 +5799,6 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_GPIO_TO_IRQ_PRESENT" "" "functions"
         ;;
 
-        migrate_vma_setup)
-            #
-            # Determine if migrate_vma_setup() function is present
-            #
-            # migrate_vma_setup() function was added by commit
-            # a7d1f22bb74f32cf3cd93f52776007e161f1a738 ("mm: turn migrate_vma
-            # upside down) in v5.4.
-            # (2019-08-20).
-            CODE="
-            #include <linux/migrate.h>
-            int conftest_migrate_vma_setup(void) {
-                migrate_vma_setup();
-            }"
-
-            compile_check_conftest "$CODE" "NV_MIGRATE_VMA_SETUP_PRESENT" "" "functions"
-        ;;
-
         migrate_vma_added_flags)
             #
             # Determine if migrate_vma structure has flags
@@ -5705,23 +5889,25 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_IOASID_GET_PRESENT" "" "functions"
         ;;
 
-        mm_pasid_set)
+        mm_pasid_drop)
             #
-            # Determine if mm_pasid_set() function is present
+            # Determine if mm_pasid_drop() function is present
+            #
+            # Added by commit 701fac40384f ("iommu/sva: Assign a PASID to mm
+            # on PASID allocation and free it on mm exit") in v5.18.
+            # Moved to linux/iommu.h in commit cd3891158a77 ("iommu/sva: Move
+            # PASID helpers to sva code") in v6.4.
             #
-            # mm_pasid_set() function was added by commit
-            # 701fac40384f07197b106136012804c3cae0b3de (iommu/sva: Assign a
-            # PASID to mm on PASID allocation and free it on mm exit) in v5.18.
-            # (2022-02-15).
             CODE="
             #if defined(NV_LINUX_SCHED_MM_H_PRESENT)
             #include <linux/sched/mm.h>
             #endif
-            void conftest_mm_pasid_set(void) {
-                mm_pasid_set();
+            #include <linux/iommu.h>
+            void conftest_mm_pasid_drop(void) {
+                mm_pasid_drop();
             }"
 
-            compile_check_conftest "$CODE" "NV_MM_PASID_SET_PRESENT" "" "functions"
+            compile_check_conftest "$CODE" "NV_MM_PASID_DROP_PRESENT" "" "functions"
         ;;
 
         drm_crtc_state_has_no_vblank)
@@ -5988,6 +6174,32 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_NUM_REGISTERED_FB_PRESENT" "" "types"
         ;;
 
+        acpi_video_register_backlight)
+            #
+            # Determine if acpi_video_register_backlight() function is present
+            #
+            # acpi_video_register_backlight was added by commit 3dbc80a3e4c55c
+            # (ACPI: video: Make backlight class device registration a separate
+            # step (v2)) for v6.0 (2022-09-02).
+            # Note: the include directive for <linux/types> in this conftest is
+            # necessary in order to support kernels between commit 0b9f7d93ca61
+            # ("ACPI / i915: ignore firmware requests backlight change") for
+            # v3.16 (2014-07-07) and commit 3bd6bce369f5 ("ACPI / video: Port
+            # to new backlight interface selection API") for v4.2 (2015-07-16).
+            # Kernels within this range use the 'bool' type and the related
+            # 'false' value in <acpi/video.h> without first including the
+            # definitions of that type and value.
+            #
+            CODE="
+            #include <linux/types.h>
+            #include <acpi/video.h>
+            void conftest_acpi_video_register_backlight(void) {
+                acpi_video_register_backlight(0);
+            }"
+
+            compile_check_conftest "$CODE" "NV_ACPI_VIDEO_REGISTER_BACKLIGHT" "" "functions"
+        ;;
+
         acpi_video_backlight_use_native)
             #
             # Determine if acpi_video_backlight_use_native() function is present
@@ -6241,23 +6453,112 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_MEMORY_FAILURE_MF_SW_SIMULATED_DEFINED" "" "types"
         ;;
 
+        drm_client_setup)
+            #
+            # Determine whether drm_client_setup is present.
+            #
+            # Added by commit d07fdf922592 ("drm/fbdev-ttm:
+            # Convert to client-setup") in v6.13.
+            #
+            CODE="
+            #include <drm/drm_fb_helper.h>
+            #if defined(NV_DRM_DRM_CLIENT_SETUP_H_PRESENT)
+            #include <drm/drm_client_setup.h>
+            #endif
+            void conftest_drm_client_setup(void) {
+                drm_client_setup();
+            }"
+
+            compile_check_conftest "$CODE" "NV_DRM_CLIENT_SETUP_PRESENT" "" "functions"
+        ;;
+
+        drm_output_poll_changed)
+            #
+            # Determine whether drm_mode_config_funcs.output_poll_changed
+            # callback is present
+            #
+            # Removed by commit 446d0f4849b1 ("drm: Remove struct
+            # drm_mode_config_funcs.output_poll_changed") in v6.12. Hotplug
+            # event support is handled through the fbdev emulation interface
+            # going forward.
+            #
+            CODE="
+            #if defined(NV_DRM_DRM_MODE_CONFIG_H_PRESENT)
+            #include <drm/drm_mode_config.h>
+            #else
+            #include <drm/drm_crtc.h>
+            #endif
+            int conftest_drm_output_poll_changed_available(void) {
+                return offsetof(struct drm_mode_config_funcs, output_poll_changed);
+            }"
+
+            compile_check_conftest "$CODE" "NV_DRM_OUTPUT_POLL_CHANGED_PRESENT" "" "types"
+        ;;
+
+        aperture_remove_conflicting_devices)
+            #
+            # Determine whether aperture_remove_conflicting_devices is present.
+            # 
+            # Added by commit 7283f862bd991 ("drm: Implement DRM aperture
+            # helpers under video/") in v6.0
+            CODE="
+            #if defined(NV_LINUX_APERTURE_H_PRESENT)
+            #include <linux/aperture.h>
+            #endif
+            void conftest_aperture_remove_conflicting_devices(void) {
+                aperture_remove_conflicting_devices();
+            }"
+            compile_check_conftest "$CODE" "NV_APERTURE_REMOVE_CONFLICTING_DEVICES_PRESENT" "" "functions"
+        ;;
+
+        aperture_remove_conflicting_pci_devices)
+            #
+            # Determine whether aperture_remove_conflicting_pci_devices is present.
+            #
+            # Added by commit 7283f862bd991 ("drm: Implement DRM aperture
+            # helpers under video/") in v6.0
+            CODE="
+            #if defined(NV_LINUX_APERTURE_H_PRESENT)
+            #include <linux/aperture.h>
+            #endif
+            void conftest_aperture_remove_conflicting_pci_devices(void) {
+                aperture_remove_conflicting_pci_devices();
+            }"
+            compile_check_conftest "$CODE" "NV_APERTURE_REMOVE_CONFLICTING_PCI_DEVICES_PRESENT" "" "functions"
+        ;;
+
+        crypto_tfm_ctx_aligned)
+            # Determine if 'crypto_tfm_ctx_aligned' is defined.
+            #
+            # Removed by commit 25c74a39e0f6 ("crypto: hmac - remove unnecessary
+            # alignment logic") in v6.7.
+            #
+            CODE="
+            #include <crypto/algapi.h>
+            void conftest_crypto_tfm_ctx_aligned(void) {
+                  (void)crypto_tfm_ctx_aligned();
+            }"
+
+            compile_check_conftest "$CODE" "NV_CRYPTO_TFM_CTX_ALIGNED_PRESENT" "" "functions"
+        ;;
+
         crypto)
             #
             # Determine if we support various crypto functions.
             # This test is not complete and may return false positive.
             #
             CODE="
-	    #include <crypto/akcipher.h>
-	    #include <crypto/algapi.h>
-	    #include <crypto/ecc_curve.h>
-	    #include <crypto/ecdh.h>
-	    #include <crypto/hash.h>
-	    #include <crypto/internal/ecc.h>
-	    #include <crypto/kpp.h>
-	    #include <crypto/public_key.h>
-	    #include <crypto/sm3.h>
-	    #include <keys/asymmetric-type.h>
-	    #include <linux/crypto.h>
+            #include <crypto/akcipher.h>
+            #include <crypto/algapi.h>
+            #include <crypto/ecc_curve.h>
+            #include <crypto/ecdh.h>
+            #include <crypto/hash.h>
+            #include <crypto/internal/ecc.h>
+            #include <crypto/kpp.h>
+            #include <crypto/public_key.h>
+            #include <crypto/sm3.h>
+            #include <keys/asymmetric-type.h>
+            #include <linux/crypto.h>
             void conftest_crypto(void) {
                 struct shash_desc sd;
                 struct crypto_shash cs;
@@ -6267,6 +6568,47 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_CRYPTO_PRESENT" "" "symbols"
         ;;
 
+        crypto_akcipher_verify)
+            #
+            # Determine whether the crypto_akcipher_verify API is still present.
+            # It was removed by commit 6b34562 ('crypto: akcipher - Drop sign/verify operations')
+            # in v6.13-rc1 (2024-10-04).
+            #
+            # This test is dependent on the crypto conftest to determine whether crypto should be
+            # enabled at all. That means that if the kernel is old enough such that crypto_akcipher_verify
+            #
+            # The test merely checks for the presence of the API, as it assumes that if the API
+            # is no longer present, the new API to replace it (crypto_sig_verify) must be present.
+            # If the kernel version is too old to have crypto_akcipher_verify, it will fail the crypto
+            # conftest above and all crypto code will be compiled out.
+            #
+            CODE="
+            #include <crypto/akcipher.h>
+            #include <linux/crypto.h>
+            void conftest_crypto_akcipher_verify(void) {
+                (void)crypto_akcipher_verify;
+            }"
+
+            compile_check_conftest "$CODE" "NV_CRYPTO_AKCIPHER_VERIFY_PRESENT" "" "symbols"
+            ;;
+
+        ecc_digits_from_bytes)
+            #
+            # Determine whether ecc_digits_from_bytes is present.
+            # It was added in commit c6ab5c915da4 ('crypto: ecc - Prevent ecc_digits_from_bytes from
+            # reading too many bytes') in v6.10.
+            #
+            # This functionality is needed when crypto_akcipher_verify is not present.
+            #
+            CODE="
+            #include <crypto/internal/ecc.h>
+            void conftest_ecc_digits_from_bytes(void) {
+                (void)ecc_digits_from_bytes;
+            }"
+
+            compile_check_conftest "$CODE" "NV_ECC_DIGITS_FROM_BYTES_PRESENT" "" "symbols"
+        ;;
+
         mempolicy_has_unified_nodes)
             #
             # Determine if the 'mempolicy' structure has
@@ -6303,6 +6645,158 @@ compile_test() {
             compile_check_conftest "$CODE" "NV_MEMPOLICY_HAS_HOME_NODE" "" "types"
         ;;
 
+        mpol_preferred_many_present)
+            #
+            # Determine if MPOL_PREFERRED_MANY enum is present or not
+            #
+            # Added by commit b27abaccf8e8b ("mm/mempolicy: add
+            # MPOL_PREFERRED_MANY for multiple preferred nodes") in
+            # v5.15
+            #
+            CODE="
+            #include <linux/mempolicy.h>
+            int mpol_preferred_many = MPOL_PREFERRED_MANY;
+            "
+
+            compile_check_conftest "$CODE" "NV_MPOL_PREFERRED_MANY_PRESENT" "" "types"
+        ;;
+
+        mmu_interval_notifier)
+            #
+            # Determine if mmu_interval_notifier struct is present or not
+            #
+            # Added by commit 99cb252f5 ("mm/mmu_notifier: add an interval tree
+            # notifier") in v5.10 (2019-11-12).
+            #
+            CODE="
+            #include <linux/mmu_notifier.h>
+            struct mmu_interval_notifier interval_notifier;
+            "
+
+            compile_check_conftest "$CODE" "NV_MMU_INTERVAL_NOTIFIER" "" "types"
+        ;;
+
+        drm_unlocked_ioctl_flag_present)
+            # Determine if DRM_UNLOCKED IOCTL flag is present.
+            #
+            # DRM_UNLOCKED was removed by commit 2798ffcc1d6a ("drm: Remove
+            # locking for legacy ioctls and DRM_UNLOCKED") in Linux
+            # next-20231208.
+            #
+            # DRM_UNLOCKED definition was moved from drmP.h to drm_ioctl.h by
+            # commit 2640981f3600 ("drm: document drm_ioctl.[hc]") in v4.12.
+            CODE="
+            #if defined(NV_DRM_DRM_IOCTL_H_PRESENT)
+            #include <drm/drm_ioctl.h>
+            #endif
+            #if defined(NV_DRM_DRMP_H_PRESENT)
+            #include <drm/drmP.h>
+            #endif
+            int flags = DRM_UNLOCKED;"
+
+            compile_check_conftest "$CODE" "NV_DRM_UNLOCKED_IOCTL_FLAG_PRESENT" "" "types"
+        ;;
+
+        page_pgmap)
+            #
+            # Determine if the page_pgmap() function is present.
+            #
+            # Added by commit 82ba975e4c43 ("mm: allow compound zone device
+            # pages") in v6.14
+            #
+            CODE="
+            #include <linux/mmzone.h>
+            int conftest_page_pgmap(void) {
+                return page_pgmap();
+            }"
+
+            compile_check_conftest "$CODE" "NV_PAGE_PGMAP_PRESENT" "" "functions"
+        ;;
+
+    folio_test_swapcache)
+            #
+            # Determine if the folio_test_swapcache() function is present.
+            #
+            # folio_test_swapcache() was exported by commit d389a4a811551 ("mm:
+            # Add folio flag manipulation functions") in v5.16.
+            #
+            CODE="
+            #include <linux/page-flags.h>
+            void conftest_folio_test_swapcache(void) {
+                folio_test_swapcache();
+            }"
+
+            compile_check_conftest "$CODE" "NV_FOLIO_TEST_SWAPCACHE_PRESENT" "" "functions"
+        ;;
+
+        module_import_ns_takes_constant)
+            #
+            # Determine if the MODULE_IMPORT_NS macro takes a string literal
+            # or constant.
+            #
+            # Commit cdd30ebb1b9f ("module: Convert symbol namespace to
+            # string literal") changed MODULE_IMPORT_NS to take a string
+            # literal in Linux kernel v6.13.
+            #
+            CODE="
+            #include <linux/module.h>
+
+            MODULE_IMPORT_NS(DMA_BUF);"
+
+            compile_check_conftest "$CODE" "NV_MODULE_IMPORT_NS_TAKES_CONSTANT" "" "generic"
+        ;;
+
+        drm_driver_has_date)
+            #
+            # Determine if the 'drm_driver' structure has a 'date' field.
+            #
+            # Removed by commit cb2e1c2136f7 ("drm: remove driver date from
+            # struct drm_driver and all drivers") in linux-next, expected in
+            # v6.14.
+            #
+            CODE="
+            #if defined(NV_DRM_DRMP_H_PRESENT)
+            #include <drm/drmP.h>
+            #endif
+
+            #if defined(NV_DRM_DRM_DRV_H_PRESENT)
+            #include <drm/drm_drv.h>
+            #endif
+
+            int conftest_drm_driver_has_date(void) {
+                return offsetof(struct drm_driver, date);
+            }"
+
+            compile_check_conftest "$CODE" "NV_DRM_DRIVER_HAS_DATE" "" "types"
+        ;;
+
+        drm_connector_helper_funcs_mode_valid_has_const_mode_arg)
+            #
+            # Determine if the 'mode' pointer argument is const in
+            # drm_connector_helper_funcs::mode_valid.
+            #
+            # The 'mode' pointer argument in
+            # drm_connector_helper_funcs::mode_valid was made const by commit
+            # 26d6fd81916e ("drm/connector: make mode_valid take a const struct
+            # drm_display_mode") in linux-next, expected in v6.15.
+            #
+            CODE="
+            #if defined(NV_DRM_DRM_ATOMIC_HELPER_H_PRESENT)
+            #include <drm/drm_atomic_helper.h>
+            #endif
+
+            static int conftest_drm_connector_mode_valid(struct drm_connector *connector,
+                                                         const struct drm_display_mode *mode) {
+                return 0;
+            }
+
+            const struct drm_connector_helper_funcs conftest_drm_connector_helper_funcs = {
+                .mode_valid = conftest_drm_connector_mode_valid,
+            };"
+
+            compile_check_conftest "$CODE" "NV_DRM_CONNECTOR_HELPER_FUNCS_MODE_VALID_HAS_CONST_MODE_ARG" "" "types"
+        ;;
+
         # When adding a new conftest entry, please use the correct format for
         # specifying the relevant upstream Linux kernel commit.
         #
@@ -6602,10 +7096,12 @@ case "$5" in
         #
         VERBOSE=$6
         iommu=CONFIG_VFIO_IOMMU_TYPE1
+        iommufd_vfio_container=CONFIG_IOMMUFD_VFIO_CONTAINER
         mdev=CONFIG_VFIO_MDEV
         kvm=CONFIG_KVM_VFIO
         vfio_pci_core=CONFIG_VFIO_PCI_CORE
         VFIO_IOMMU_PRESENT=0
+        VFIO_IOMMUFD_VFIO_CONTAINER_PRESENT=0
         VFIO_MDEV_PRESENT=0
         KVM_PRESENT=0
         VFIO_PCI_CORE_PRESENT=0
@@ -6615,6 +7111,10 @@ case "$5" in
                 VFIO_IOMMU_PRESENT=1
             fi
 
+            if (test_configuration_option ${iommufd_vfio_container} || test_configuration_option ${iommufd_vfio_container}_MODULE); then
+                VFIO_IOMMUFD_VFIO_CONTAINER_PRESENT=1
+            fi
+
             if (test_configuration_option ${mdev} || test_configuration_option ${mdev}_MODULE); then
                 VFIO_MDEV_PRESENT=1
             fi
@@ -6627,34 +7127,21 @@ case "$5" in
                 VFIO_PCI_CORE_PRESENT=1
             fi
 
-            # When this sanity check is run via nvidia-installer, it sets ARCH as aarch64.
-            # But, when it is run via Kbuild, ARCH is set as arm64
-            if [ "$ARCH" = "aarch64" ]; then
-                ARCH="arm64"
-            fi
-
-            if [ "$VFIO_IOMMU_PRESENT" != "0" ] && [ "$KVM_PRESENT" != "0" ] ; then
-
-                # On x86_64, vGPU requires MDEV framework to be present.
-                # On aarch64, vGPU requires vfio-pci-core framework to be present.
-                if ([ "$ARCH" = "arm64" ] && [ "$VFIO_PCI_CORE_PRESENT" != "0" ]) ||
-                   ([ "$ARCH" = "x86_64" ] && [ "$VFIO_MDEV_PRESENT" != "0" ];) then
+            if ([ "$VFIO_IOMMU_PRESENT" != "0" ] || [ "$VFIO_IOMMUFD_VFIO_CONTAINER_PRESENT" != "0" ])&& [ "$KVM_PRESENT" != "0" ] ; then
+                # vGPU requires either MDEV or vfio-pci-core framework to be present.
+                if [ "$VFIO_MDEV_PRESENT" != "0" ] || [ "$VFIO_PCI_CORE_PRESENT" != "0" ]; then
                     exit 0
                 fi
             fi
 
             echo "Below CONFIG options are missing on the kernel for installing";
             echo "NVIDIA vGPU driver on KVM host";
-            if [ "$VFIO_IOMMU_PRESENT" = "0" ]; then
-                echo "CONFIG_VFIO_IOMMU_TYPE1";
+            if [ "$VFIO_IOMMU_PRESENT" = "0" ] && [ "$VFIO_IOMMUFD_VFIO_CONTAINER_PRESENT" = "0" ]; then
+                echo "either CONFIG_VFIO_IOMMU_TYPE1 or CONFIG_IOMMUFD_VFIO_CONTAINER";
             fi
 
-            if [ "$ARCH" = "arm64" ] &&  [ "$VFIO_PCI_CORE_PRESENT" = "0" ]; then
-                echo "CONFIG_VFIO_PCI_CORE";
-            fi
-
-            if [ "$ARCH" = "x86_64" ] && [ "$VFIO_MDEV_PRESENT" = "0" ]; then
-                echo "CONFIG_VFIO_MDEV";
+            if [ "$VFIO_MDEV_PRESENT" = "0" ] && [ "$VFIO_PCI_CORE_PRESENT" = "0" ]; then
+                echo "either CONFIG_VFIO_MDEV or CONFIG_VFIO_PCI_CORE";
             fi
 
             if [ "$KVM_PRESENT" = "0" ]; then

kernel-open/nvidia-drm/nvidia-drm-connector.c

@@ -314,7 +314,11 @@ static int nv_drm_connector_get_modes(struct drm_connector *connector)
 }
 
 static int nv_drm_connector_mode_valid(struct drm_connector    *connector,
+#if defined(NV_DRM_CONNECTOR_HELPER_FUNCS_MODE_VALID_HAS_CONST_MODE_ARG)
+                                       const struct drm_display_mode *mode)
+#else
                                        struct drm_display_mode *mode)
+#endif
 {
     struct drm_device *dev = connector->dev;
     struct nv_drm_device *nv_dev = to_nv_device(dev);

kernel-open/nvidia-drm/nvidia-drm-drv.c

@@ -105,6 +105,7 @@ static const char* nv_get_input_colorspace_name(
 
 #if defined(NV_DRM_ATOMIC_MODESET_AVAILABLE)
 
+#if defined(NV_DRM_OUTPUT_POLL_CHANGED_PRESENT)
 static void nv_drm_output_poll_changed(struct drm_device *dev)
 {
     struct drm_connector *connector = NULL;
@@ -148,15 +149,19 @@ static void nv_drm_output_poll_changed(struct drm_device *dev)
     nv_drm_connector_list_iter_end(&conn_iter);
 #endif
 }
+#endif /* NV_DRM_OUTPUT_POLL_CHANGED_PRESENT */
 
 static struct drm_framebuffer *nv_drm_framebuffer_create(
     struct drm_device *dev,
     struct drm_file *file,
-    #if defined(NV_DRM_HELPER_MODE_FILL_FB_STRUCT_HAS_CONST_MODE_CMD_ARG)
+#if defined(NV_DRM_FB_CREATE_TAKES_FORMAT_INFO)
+    const struct drm_format_info *info,
+#endif
+#if defined(NV_DRM_HELPER_MODE_FILL_FB_STRUCT_HAS_CONST_MODE_CMD_ARG)
     const struct drm_mode_fb_cmd2 *cmd
-    #else
+#else
     struct drm_mode_fb_cmd2 *cmd
-    #endif
+#endif
 )
 {
     struct drm_mode_fb_cmd2 local_cmd;
@@ -167,11 +172,14 @@ static struct drm_framebuffer *nv_drm_framebuffer_create(
     fb = nv_drm_internal_framebuffer_create(
             dev,
             file,
+#if defined(NV_DRM_FB_CREATE_TAKES_FORMAT_INFO)
+            info,
+#endif
             &local_cmd);
 
-    #if !defined(NV_DRM_HELPER_MODE_FILL_FB_STRUCT_HAS_CONST_MODE_CMD_ARG)
+#if !defined(NV_DRM_HELPER_MODE_FILL_FB_STRUCT_HAS_CONST_MODE_CMD_ARG)
     *cmd = local_cmd;
-    #endif
+#endif
 
     return fb;
 }
@@ -185,7 +193,9 @@ static const struct drm_mode_config_funcs nv_mode_config_funcs = {
     .atomic_check  = nv_drm_atomic_check,
     .atomic_commit = nv_drm_atomic_commit,
 
+    #if defined(NV_DRM_OUTPUT_POLL_CHANGED_PRESENT)
     .output_poll_changed = nv_drm_output_poll_changed,
+    #endif
 };
 
 static void nv_drm_event_callback(const struct NvKmsKapiEvent *event)
@@ -1281,6 +1291,10 @@ static const struct file_operations nv_drm_fops = {
     .read           = drm_read,
 
     .llseek         = noop_llseek,
+
+#if defined(NV_FILE_OPERATIONS_FOP_UNSIGNED_OFFSET_PRESENT)
+    .fop_flags   = FOP_UNSIGNED_OFFSET,
+#endif
 };
 
 static const struct drm_ioctl_desc nv_drm_ioctls[] = {
@@ -1312,9 +1326,21 @@ static const struct drm_ioctl_desc nv_drm_ioctls[] = {
                       DRM_RENDER_ALLOW|DRM_UNLOCKED),
 #endif
 
+    /*
+     * DRM_UNLOCKED is implicit for all non-legacy DRM driver IOCTLs since Linux
+     * v4.10 commit fa5386459f06 "drm: Used DRM_LEGACY for all legacy functions"
+     * (Linux v4.4 commit ea487835e887 "drm: Enforce unlocked ioctl operation
+     * for kms driver ioctls" previously did it only for drivers that set the
+     * DRM_MODESET flag), so this will race with SET_CLIENT_CAP. Linux v4.11
+     * commit dcf727ab5d17 "drm: setclientcap doesn't need the drm BKL" also
+     * removed locking from SET_CLIENT_CAP so there is no use attempting to lock
+     * manually. The latter commit acknowledges that this can expose userspace
+     * to inconsistent behavior when racing with itself, but accepts that risk.
+     */
     DRM_IOCTL_DEF_DRV(NVIDIA_GET_CLIENT_CAPABILITY,
                       nv_drm_get_client_capability_ioctl,
                       0),
+
 #if defined(NV_DRM_ATOMIC_MODESET_AVAILABLE)
     DRM_IOCTL_DEF_DRV(NVIDIA_GET_CRTC_CRC32,
                       nv_drm_get_crtc_crc32_ioctl,
@@ -1367,8 +1393,23 @@ static struct drm_driver nv_drm_driver = {
     .ioctls                 = nv_drm_ioctls,
     .num_ioctls             = ARRAY_SIZE(nv_drm_ioctls),
 
+/*
+ * linux-next commit 71a7974ac701 ("drm/prime: Unexport helpers for fd/handle
+ * conversion") unexports drm_gem_prime_handle_to_fd() and
+ * drm_gem_prime_fd_to_handle().
+ *
+ * Prior linux-next commit 6b85aa68d9d5 ("drm: Enable PRIME import/export for
+ * all drivers") made these helpers the default when .prime_handle_to_fd /
+ * .prime_fd_to_handle are unspecified, so it's fine to just skip specifying
+ * them if the helpers aren't present.
+ */
+#if NV_IS_EXPORT_SYMBOL_PRESENT_drm_gem_prime_handle_to_fd
     .prime_handle_to_fd     = drm_gem_prime_handle_to_fd,
+#endif
+#if NV_IS_EXPORT_SYMBOL_PRESENT_drm_gem_prime_fd_to_handle
     .prime_fd_to_handle     = drm_gem_prime_fd_to_handle,
+#endif
+
     .gem_prime_import       = nv_drm_gem_prime_import,
     .gem_prime_import_sg_table = nv_drm_gem_prime_import_sg_table,
 
@@ -1404,7 +1445,10 @@ static struct drm_driver nv_drm_driver = {
     .name                   = "nvidia-drm",
 
     .desc                   = "NVIDIA DRM driver",
+
+#if defined(NV_DRM_DRIVER_HAS_DATE)
     .date                   = "20160202",
+#endif
 
 #if defined(NV_DRM_DRIVER_HAS_DEVICE_LIST)
     .device_list            = LIST_HEAD_INIT(nv_drm_driver.device_list),

kernel-open/nvidia-drm/nvidia-drm-fb.c

@@ -206,6 +206,9 @@ fail:
 struct drm_framebuffer *nv_drm_internal_framebuffer_create(
     struct drm_device *dev,
     struct drm_file *file,
+#if defined(NV_DRM_FB_CREATE_TAKES_FORMAT_INFO)
+    const struct drm_format_info *info,
+#endif
     struct drm_mode_fb_cmd2 *cmd)
 {
     struct nv_drm_device *nv_dev = to_nv_device(dev);
@@ -259,6 +262,9 @@ struct drm_framebuffer *nv_drm_internal_framebuffer_create(
         dev,
         #endif
         &nv_fb->base,
+#if defined(NV_DRM_FB_CREATE_TAKES_FORMAT_INFO)
+        info,
+#endif 
         cmd);
 
     /*

kernel-open/nvidia-drm/nvidia-drm-fb.h

@@ -59,6 +59,9 @@ static inline struct nv_drm_framebuffer *to_nv_framebuffer(
 struct drm_framebuffer *nv_drm_internal_framebuffer_create(
     struct drm_device *dev,
     struct drm_file *file,
+#if defined(NV_DRM_FB_CREATE_TAKES_FORMAT_INFO)
+    const struct drm_format_info *info,
+#endif
     struct drm_mode_fb_cmd2 *cmd);
 
 #endif /* NV_DRM_ATOMIC_MODESET_AVAILABLE */

kernel-open/nvidia-drm/nvidia-drm-gem-nvkms-memory.c

@@ -243,6 +243,15 @@ static int __nv_drm_nvkms_gem_obj_init(
     NvU64 *pages = NULL;
     NvU32 numPages = 0;
 
+    if ((size % PAGE_SIZE) != 0) {
+        NV_DRM_DEV_LOG_ERR(
+            nv_dev,
+            "NvKmsKapiMemory 0x%p size should be in a multiple of page size to "
+            "create a gem object",
+            pMemory);
+        return -EINVAL;
+    }
+
     nv_nvkms_memory->pPhysicalAddress = NULL;
     nv_nvkms_memory->pWriteCombinedIORemapAddress = NULL;
     nv_nvkms_memory->physically_mapped = false;

kernel-open/nvidia-drm/nvidia-drm-helper.h

@@ -582,6 +582,19 @@ static inline int nv_drm_format_num_planes(uint32_t format)
 
 #endif /* defined(NV_DRM_FORMAT_MODIFIERS_PRESENT) */
 
+/*
+ * DRM_UNLOCKED was removed with linux-next commit 2798ffcc1d6a ("drm: Remove
+ * locking for legacy ioctls and DRM_UNLOCKED"), but it was previously made
+ * implicit for all non-legacy DRM driver IOCTLs since Linux v4.10 commit
+ * fa5386459f06 "drm: Used DRM_LEGACY for all legacy functions" (Linux v4.4
+ * commit ea487835e887 "drm: Enforce unlocked ioctl operation for kms driver
+ * ioctls" previously did it only for drivers that set the DRM_MODESET flag), so
+ * it was effectively a no-op anyway.
+ */
+#if !defined(NV_DRM_UNLOCKED_IOCTL_FLAG_PRESENT)
+#define DRM_UNLOCKED 0
+#endif
+
 /*
  * drm_vma_offset_exact_lookup_locked() were added
  * by kernel commit 2225cfe46bcc which was Signed-off-by:

kernel-open/nvidia-drm/nvidia-drm-modeset.c

@@ -451,6 +451,13 @@ int nv_drm_atomic_commit(struct drm_device *dev,
 #else
     drm_atomic_helper_swap_state(dev, state);
 #endif
+    /*
+     * Used to update legacy modeset state pointers to support UAPIs not updated
+     * by the core atomic modeset infrastructure.
+     *
+     * Example: /sys/class/drm/<card connector>/enabled
+     */
+    drm_atomic_helper_update_legacy_modeset_state(dev, state);
 
     /*
      * nv_drm_atomic_commit_internal() must not return failure after

kernel-open/nvidia-drm/nvidia-drm.Kbuild

@@ -54,6 +54,9 @@ NV_CONFTEST_GENERIC_COMPILE_TESTS += drm_atomic_available
 NV_CONFTEST_GENERIC_COMPILE_TESTS += is_export_symbol_gpl_refcount_inc
 NV_CONFTEST_GENERIC_COMPILE_TESTS += is_export_symbol_gpl_refcount_dec_and_test
 NV_CONFTEST_GENERIC_COMPILE_TESTS += drm_alpha_blending_available
+NV_CONFTEST_GENERIC_COMPILE_TESTS += is_export_symbol_present_drm_gem_prime_fd_to_handle
+NV_CONFTEST_GENERIC_COMPILE_TESTS += is_export_symbol_present_drm_gem_prime_handle_to_fd
+NV_CONFTEST_GENERIC_COMPILE_TESTS += is_export_symbol_gpl___vma_start_write
 
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += drm_dev_unref
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += drm_reinit_primary_mode_group
@@ -131,3 +134,9 @@ NV_CONFTEST_TYPE_COMPILE_TESTS += drm_connector_lookup
 NV_CONFTEST_TYPE_COMPILE_TESTS += drm_connector_put
 NV_CONFTEST_TYPE_COMPILE_TESTS += vm_area_struct_has_const_vm_flags
 NV_CONFTEST_TYPE_COMPILE_TESTS += drm_driver_has_dumb_destroy
+NV_CONFTEST_TYPE_COMPILE_TESTS += drm_unlocked_ioctl_flag_present
+NV_CONFTEST_TYPE_COMPILE_TESTS += drm_output_poll_changed
+NV_CONFTEST_TYPE_COMPILE_TESTS += drm_driver_has_date
+NV_CONFTEST_TYPE_COMPILE_TESTS += file_operations_fop_unsigned_offset_present
+NV_CONFTEST_TYPE_COMPILE_TESTS += drm_connector_helper_funcs_mode_valid_has_const_mode_arg
+NV_CONFTEST_TYPE_COMPILE_TESTS += drm_fb_create_takes_format_info

kernel-open/nvidia-modeset/nv-kthread-q.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2016 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2016-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -176,7 +176,7 @@ static struct task_struct *thread_create_on_node(int (*threadfn)(void *data),
 {
 
     unsigned i, j;
-    const static unsigned attempts = 3;
+    static const unsigned attempts = 3;
     struct task_struct *thread[3];
 
     for (i = 0;; i++) {

kernel-open/nvidia-modeset/nvidia-modeset-linux.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2015-21 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2015-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -54,7 +54,11 @@
 #include "nv-time.h"
 #include "nv-lock.h"
 
-#if !defined(CONFIG_RETPOLINE)
+/*
+ * Commit aefb2f2e619b ("x86/bugs: Rename CONFIG_RETPOLINE =>
+ * CONFIG_MITIGATION_RETPOLINE) in v6.8 renamed CONFIG_RETPOLINE.
+ */
+#if !defined(CONFIG_RETPOLINE) && !defined(CONFIG_MITIGATION_RETPOLINE)
 #include "nv-retpoline.h"
 #endif
 
@@ -68,6 +72,9 @@ module_param_named(output_rounding_fix, output_rounding_fix, bool, 0400);
 static bool disable_vrr_memclk_switch = false;
 module_param_named(disable_vrr_memclk_switch, disable_vrr_memclk_switch, bool, 0400);
 
+static bool opportunistic_display_sync = true;
+module_param_named(opportunistic_display_sync, opportunistic_display_sync, bool, 0400);
+
 /* These parameters are used for fault injection tests.  Normally the defaults
  * should be used. */
 MODULE_PARM_DESC(fail_malloc, "Fail the Nth call to nvkms_alloc");
@@ -99,6 +106,11 @@ NvBool nvkms_disable_vrr_memclk_switch(void)
     return disable_vrr_memclk_switch;
 }
 
+NvBool nvkms_opportunistic_display_sync(void)
+{
+    return opportunistic_display_sync;
+}
+
 #define NVKMS_SYNCPT_STUBS_NEEDED
 
 /*************************************************************************
@@ -200,9 +212,23 @@ static inline int nvkms_read_trylock_pm_lock(void)
 
 static inline void nvkms_read_lock_pm_lock(void)
 {
-    while (!down_read_trylock(&nvkms_pm_lock)) {
-        try_to_freeze();
-        cond_resched();
+    if ((current->flags & PF_NOFREEZE)) {
+        /*
+         * Non-freezable tasks (i.e. kthreads in this case) don't have to worry
+         * about being frozen during system suspend, but do need to block so
+         * that the CPU can go idle during s2idle. Do a normal uninterruptible
+         * blocking wait for the PM lock.
+         */
+        down_read(&nvkms_pm_lock);
+    } else {
+        /*
+         * For freezable tasks, make sure we give the kernel an opportunity to
+         * freeze if taking the PM lock fails.
+         */
+        while (!down_read_trylock(&nvkms_pm_lock)) {
+            try_to_freeze();
+            cond_resched();
+        }
     }
 }
 
@@ -626,7 +652,11 @@ static void nvkms_kthread_q_callback(void *arg)
      * pending timers and than waiting for workqueue callbacks.
      */
     if (timer->kernel_timer_created) {
+#if !defined(NV_BSD) && NV_IS_EXPORT_SYMBOL_PRESENT_timer_delete_sync
+        timer_delete_sync(&timer->kernel_timer);
+#else
         del_timer_sync(&timer->kernel_timer);
+#endif
     }
 
     /*
@@ -981,6 +1011,11 @@ nvkms_register_backlight(NvU32 gpu_id, NvU32 display_id, void *drv_priv,
 
 #if defined(NV_ACPI_VIDEO_BACKLIGHT_USE_NATIVE)
     if (!acpi_video_backlight_use_native()) {
+#if defined(NV_ACPI_VIDEO_REGISTER_BACKLIGHT)
+        nvkms_log(NVKMS_LOG_LEVEL_INFO, NVKMS_LOG_PREFIX,
+                  "ACPI reported no NVIDIA native backlight available; attempting to use ACPI backlight.");
+        acpi_video_register_backlight();
+#endif
         return NULL;
     }
 #endif
@@ -1055,7 +1090,7 @@ static void nvkms_kapi_event_kthread_q_callback(void *arg)
     nvKmsKapiHandleEventQueueChange(device);
 }
 
-struct nvkms_per_open *nvkms_open_common(enum NvKmsClientType type,
+static struct nvkms_per_open *nvkms_open_common(enum NvKmsClientType type,
                                          struct NvKmsKapiDevice *device,
                                          int *status)
 {
@@ -1107,7 +1142,7 @@ failed:
     return NULL;
 }
 
-void nvkms_close_pm_locked(struct nvkms_per_open *popen)
+static void nvkms_close_pm_locked(struct nvkms_per_open *popen)
 {
     /*
      * Don't use down_interruptible(): we need to free resources
@@ -1170,7 +1205,7 @@ static void nvkms_close_popen(struct nvkms_per_open *popen)
     }
 }
 
-int nvkms_ioctl_common
+static int nvkms_ioctl_common
 (
     struct nvkms_per_open *popen,
     NvU32 cmd, NvU64 address, const size_t size
@@ -1712,7 +1747,11 @@ restart:
              * completion, and we wait for queue completion with
              * nv_kthread_q_stop below.
              */
+#if !defined(NV_BSD) && NV_IS_EXPORT_SYMBOL_PRESENT_timer_delete_sync
+            if (timer_delete_sync(&timer->kernel_timer) == 1) {
+#else
             if (del_timer_sync(&timer->kernel_timer) == 1) {
+#endif
                 /*  We've deactivated timer so we need to clean after it */
                 list_del(&timer->timers_list);
 

kernel-open/nvidia-modeset/nvidia-modeset-os-interface.h

@@ -99,6 +99,7 @@ typedef struct {
 NvBool nvkms_output_rounding_fix(void);
 
 NvBool nvkms_disable_vrr_memclk_switch(void);
+NvBool nvkms_opportunistic_display_sync(void);
 
 void   nvkms_call_rm    (void *ops);
 void*  nvkms_alloc      (size_t size,

kernel-open/nvidia-modeset/nvidia-modeset.Kbuild

@@ -40,9 +40,6 @@ NV_KERNEL_MODULE_TARGETS += $(NVIDIA_MODESET_KO)
 NVIDIA_MODESET_BINARY_OBJECT := $(src)/nvidia-modeset/nv-modeset-kernel.o_binary
 NVIDIA_MODESET_BINARY_OBJECT_O := nvidia-modeset/nv-modeset-kernel.o
 
-quiet_cmd_symlink = SYMLINK $@
-cmd_symlink = ln -sf $< $@
-
 targets += $(NVIDIA_MODESET_BINARY_OBJECT_O)
 
 $(obj)/$(NVIDIA_MODESET_BINARY_OBJECT_O): $(NVIDIA_MODESET_BINARY_OBJECT) FORCE
@@ -93,3 +90,5 @@ NV_CONFTEST_FUNCTION_COMPILE_TESTS += list_is_first
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += ktime_get_real_ts64
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += ktime_get_raw_ts64
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += acpi_video_backlight_use_native
+NV_CONFTEST_FUNCTION_COMPILE_TESTS += acpi_video_register_backlight
+NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_timer_delete_sync

kernel-open/nvidia-peermem/nvidia-peermem.c

@@ -1,20 +1,25 @@
-/* SPDX-License-Identifier: Linux-OpenIB */
 /*
  * Copyright (c) 2006, 2007 Cisco Systems, Inc. All rights reserved.
  * Copyright (c) 2007, 2008 Mellanox Technologies. All rights reserved.
  *
- * Redistribution and use in source and binary forms, with or
- * without modification, are permitted provided that the following
- * conditions are met:
+ * This software is available to you under a choice of one of two
+ * licenses.  You may choose to be licensed under the terms of the GNU
+ * General Public License (GPL) Version 2, available from the file
+ * COPYING in the main directory of this source tree, or the
+ * OpenIB.org BSD license below:
  *
- *  - Redistributions of source code must retain the above
- *    copyright notice, this list of conditions and the following
- *    disclaimer.
+ *     Redistribution and use in source and binary forms, with or
+ *     without modification, are permitted provided that the following
+ *     conditions are met:
  *
- *  - Redistributions in binary form must reproduce the above
- *    copyright notice, this list of conditions and the following
- *    disclaimer in the documentation and/or other materials
- *    provided with the distribution.
+ *      - Redistributions of source code must retain the above
+ *        copyright notice, this list of conditions and the following
+ *        disclaimer.
+ *
+ *      - Redistributions in binary form must reproduce the above
+ *        copyright notice, this list of conditions and the following
+ *        disclaimer in the documentation and/or other materials
+ *        provided with the distribution.
  *
  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
@@ -43,7 +48,9 @@
 
 MODULE_AUTHOR("Yishai Hadas");
 MODULE_DESCRIPTION("NVIDIA GPU memory plug-in");
-MODULE_LICENSE("Linux-OpenIB");
+
+MODULE_LICENSE("Dual BSD/GPL");
+
 MODULE_VERSION(DRV_VERSION);
 enum {
         NV_MEM_PEERDIRECT_SUPPORT_DEFAULT = 0,
@@ -53,7 +60,13 @@ static int peerdirect_support = NV_MEM_PEERDIRECT_SUPPORT_DEFAULT;
 module_param(peerdirect_support, int, S_IRUGO);
 MODULE_PARM_DESC(peerdirect_support, "Set level of support for Peer-direct, 0 [default] or 1 [legacy, for example MLNX_OFED 4.9 LTS]");
 
-#define peer_err(FMT, ARGS...) printk(KERN_ERR "nvidia-peermem" " %s:%d " FMT, __FUNCTION__, __LINE__, ## ARGS)
+
+#define peer_err(FMT, ARGS...) printk(KERN_ERR "nvidia-peermem" " %s:%d ERROR " FMT, __FUNCTION__, __LINE__, ## ARGS)
+#ifdef NV_MEM_DEBUG
+#define peer_trace(FMT, ARGS...) printk(KERN_DEBUG "nvidia-peermem" " %s:%d TRACE " FMT, __FUNCTION__, __LINE__, ## ARGS)
+#else
+#define peer_trace(FMT, ARGS...) do {} while (0)
+#endif
 
 #if defined(NV_MLNX_IB_PEER_MEM_SYMBOLS_PRESENT)
 
@@ -74,7 +87,10 @@ invalidate_peer_memory mem_invalidate_callback;
 static void *reg_handle = NULL;
 static void *reg_handle_nc = NULL;
 
+#define NV_MEM_CONTEXT_MAGIC ((u64)0xF1F4F1D0FEF0DAD0ULL)
+
 struct nv_mem_context {
+    u64 pad1;
     struct nvidia_p2p_page_table *page_table;
     struct nvidia_p2p_dma_mapping *dma_mapping;
     u64 core_context;
@@ -86,8 +102,22 @@ struct nv_mem_context {
     struct task_struct *callback_task;
     int sg_allocated;
     struct sg_table sg_head;
+    u64 pad2;
 };
 
+#define NV_MEM_CONTEXT_CHECK_OK(MC) ({                                  \
+    struct nv_mem_context *mc = (MC);                                   \
+    int rc = ((0 != mc) &&                                              \
+              (READ_ONCE(mc->pad1) == NV_MEM_CONTEXT_MAGIC) &&          \
+              (READ_ONCE(mc->pad2) == NV_MEM_CONTEXT_MAGIC));           \
+    if (!rc) {                                                          \
+        peer_trace("invalid nv_mem_context=%px pad1=%016llx pad2=%016llx\n", \
+                   mc,                                                  \
+                   mc?mc->pad1:0,                                       \
+                   mc?mc->pad2:0);                                      \
+    }                                                                   \
+    rc;                                                                 \
+})
 
 static void nv_get_p2p_free_callback(void *data)
 {
@@ -97,8 +127,9 @@ static void nv_get_p2p_free_callback(void *data)
     struct nvidia_p2p_dma_mapping *dma_mapping = NULL;
 
     __module_get(THIS_MODULE);
-    if (!nv_mem_context) {
-        peer_err("nv_get_p2p_free_callback -- invalid nv_mem_context\n");
+
+    if (!NV_MEM_CONTEXT_CHECK_OK(nv_mem_context)) {
+        peer_err("detected invalid context, skipping further processing\n");
         goto out;
     }
 
@@ -169,9 +200,11 @@ static int nv_mem_acquire(unsigned long addr, size_t size, void *peer_mem_privat
         /* Error case handled as not mine */
         return 0;
 
+    nv_mem_context->pad1 = NV_MEM_CONTEXT_MAGIC;
     nv_mem_context->page_virt_start = addr & GPU_PAGE_MASK;
     nv_mem_context->page_virt_end   = (addr + size + GPU_PAGE_SIZE - 1) & GPU_PAGE_MASK;
     nv_mem_context->mapped_size  = nv_mem_context->page_virt_end - nv_mem_context->page_virt_start;
+    nv_mem_context->pad2 = NV_MEM_CONTEXT_MAGIC;
 
     ret = nvidia_p2p_get_pages(0, 0, nv_mem_context->page_virt_start, nv_mem_context->mapped_size,
                                &nv_mem_context->page_table, nv_mem_dummy_callback, nv_mem_context);
@@ -195,6 +228,7 @@ static int nv_mem_acquire(unsigned long addr, size_t size, void *peer_mem_privat
     return 1;
 
 err:
+    memset(nv_mem_context, 0, sizeof(*nv_mem_context));
     kfree(nv_mem_context);
 
     /* Error case handled as not mine */
@@ -342,6 +376,7 @@ static void nv_mem_release(void *context)
         sg_free_table(&nv_mem_context->sg_head);
         nv_mem_context->sg_allocated = 0;
     }
+    memset(nv_mem_context, 0, sizeof(*nv_mem_context));
     kfree(nv_mem_context);
     module_put(THIS_MODULE);
     return;

kernel-open/nvidia-uvm/nv-kthread-q-selftest.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016 NVIDIA Corporation
+    Copyright (c) 2016-2024 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -81,7 +81,7 @@
 #define NUM_Q_ITEMS_IN_MULTITHREAD_TEST (NUM_TEST_Q_ITEMS * NUM_TEST_KTHREADS)
 
 // This exists in order to have a function to place a breakpoint on:
-void on_nvq_assert(void)
+static void on_nvq_assert(void)
 {
     (void)NULL;
 }

kernel-open/nvidia-uvm/nv-kthread-q.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2016 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2016-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: MIT
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
@@ -176,7 +176,7 @@ static struct task_struct *thread_create_on_node(int (*threadfn)(void *data),
 {
 
     unsigned i, j;
-    const static unsigned attempts = 3;
+    static const unsigned attempts = 3;
     struct task_struct *thread[3];
 
     for (i = 0;; i++) {

kernel-open/nvidia-uvm/nvidia-uvm.Kbuild

@@ -81,12 +81,13 @@ NV_CONFTEST_FUNCTION_COMPILE_TESTS += set_memory_uc
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += set_pages_uc
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += ktime_get_raw_ts64
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += ioasid_get
-NV_CONFTEST_FUNCTION_COMPILE_TESTS += mm_pasid_set
-NV_CONFTEST_FUNCTION_COMPILE_TESTS += migrate_vma_setup
+NV_CONFTEST_FUNCTION_COMPILE_TESTS += mm_pasid_drop
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += mmget_not_zero
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += mmgrab
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += iommu_sva_bind_device_has_drvdata_arg
 NV_CONFTEST_FUNCTION_COMPILE_TESTS += vm_fault_to_errno
+NV_CONFTEST_FUNCTION_COMPILE_TESTS += folio_test_swapcache
+NV_CONFTEST_FUNCTION_COMPILE_TESTS += page_pgmap
 
 NV_CONFTEST_TYPE_COMPILE_TESTS += backing_dev_info
 NV_CONFTEST_TYPE_COMPILE_TESTS += mm_context_t
@@ -100,6 +101,7 @@ NV_CONFTEST_TYPE_COMPILE_TESTS += kmem_cache_has_kobj_remove_work
 NV_CONFTEST_TYPE_COMPILE_TESTS += sysfs_slab_unlink
 NV_CONFTEST_TYPE_COMPILE_TESTS += vm_fault_t
 NV_CONFTEST_TYPE_COMPILE_TESTS += mmu_notifier_ops_invalidate_range
+NV_CONFTEST_TYPE_COMPILE_TESTS += mmu_notifier_ops_arch_invalidate_secondary_tlbs
 NV_CONFTEST_TYPE_COMPILE_TESTS += proc_ops
 NV_CONFTEST_TYPE_COMPILE_TESTS += timespec64
 NV_CONFTEST_TYPE_COMPILE_TESTS += mm_has_mmap_lock
@@ -110,5 +112,10 @@ NV_CONFTEST_TYPE_COMPILE_TESTS += handle_mm_fault_has_mm_arg
 NV_CONFTEST_TYPE_COMPILE_TESTS += handle_mm_fault_has_pt_regs_arg
 NV_CONFTEST_TYPE_COMPILE_TESTS += mempolicy_has_unified_nodes
 NV_CONFTEST_TYPE_COMPILE_TESTS += mempolicy_has_home_node
+NV_CONFTEST_TYPE_COMPILE_TESTS += mpol_preferred_many_present
+NV_CONFTEST_TYPE_COMPILE_TESTS += mmu_interval_notifier
 
 NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_int_active_memcg
+NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_migrate_vma_setup
+NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present___iowrite64_lo_hi
+NV_CONFTEST_SYMBOL_COMPILE_TESTS += is_export_symbol_present_make_device_exclusive

kernel-open/nvidia-uvm/uvm.c

@@ -571,7 +571,6 @@ static void uvm_vm_open_managed_entry(struct vm_area_struct *vma)
 static void uvm_vm_close_managed(struct vm_area_struct *vma)
 {
     uvm_va_space_t *va_space = uvm_va_space_get(vma->vm_file);
-    uvm_processor_id_t gpu_id;
     bool make_zombie = false;
 
     if (current->mm != NULL)
@@ -606,12 +605,6 @@ static void uvm_vm_close_managed(struct vm_area_struct *vma)
 
     uvm_destroy_vma_managed(vma, make_zombie);
 
-    // Notify GPU address spaces that the fault buffer needs to be flushed to
-    // avoid finding stale entries that can be attributed to new VA ranges
-    // reallocated at the same address.
-    for_each_gpu_id_in_mask(gpu_id, &va_space->registered_gpu_va_spaces) {
-        uvm_processor_mask_set_atomic(&va_space->needs_fault_buffer_flush, gpu_id);
-    }
     uvm_va_space_up_write(va_space);
 
     if (current->mm != NULL)
@@ -689,6 +682,9 @@ static void uvm_vm_open_semaphore_pool(struct vm_area_struct *vma)
     // Semaphore pool vmas do not have vma wrappers, but some functions will
     // assume vm_private_data is a wrapper.
     vma->vm_private_data = NULL;
+#if defined(VM_WIPEONFORK)
+    nv_vm_flags_set(vma, VM_WIPEONFORK);
+#endif
 
     if (is_fork) {
         // If we forked, leave the parent vma alone.

kernel-open/nvidia-uvm/uvm_ada.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2021 NVIDIA Corporation
+    Copyright (c) 2021-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -94,4 +94,6 @@ void uvm_hal_ada_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     parent_gpu->map_remap_larger_page_promotion = false;
 
     parent_gpu->plc_supported = true;
+
+    parent_gpu->no_ats_range_required = false;
 }

kernel-open/nvidia-uvm/uvm_ampere.c

@@ -101,4 +101,6 @@ void uvm_hal_ampere_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
         parent_gpu->map_remap_larger_page_promotion = false;
 
     parent_gpu->plc_supported = true;
+
+    parent_gpu->no_ats_range_required = false;
 }

kernel-open/nvidia-uvm/uvm_ampere_ce.c

@@ -121,6 +121,8 @@ bool uvm_hal_ampere_ce_memcopy_is_valid_c6b5(uvm_push_t *push, uvm_gpu_address_t
         return true;
 
     if (uvm_channel_is_proxy(push->channel)) {
+        uvm_pushbuffer_t *pushbuffer;
+
         if (dst.is_virtual) {
             UVM_ERR_PRINT("Destination address of memcopy must be physical, not virtual\n");
             return false;
@@ -142,7 +144,8 @@ bool uvm_hal_ampere_ce_memcopy_is_valid_c6b5(uvm_push_t *push, uvm_gpu_address_t
             return false;
         }
 
-        push_begin_gpu_va = uvm_pushbuffer_get_gpu_va_for_push(push->channel->pool->manager->pushbuffer, push);
+        pushbuffer = uvm_channel_get_pushbuffer(push->channel);
+        push_begin_gpu_va = uvm_pushbuffer_get_gpu_va_for_push(pushbuffer, push);
 
         if ((src.address < push_begin_gpu_va) || (src.address >= push_begin_gpu_va + uvm_push_get_size(push))) {
             UVM_ERR_PRINT("Source address of memcopy must point to pushbuffer\n");
@@ -177,10 +180,13 @@ bool uvm_hal_ampere_ce_memcopy_is_valid_c6b5(uvm_push_t *push, uvm_gpu_address_t
 // irrespective of the virtualization mode.
 void uvm_hal_ampere_ce_memcopy_patch_src_c6b5(uvm_push_t *push, uvm_gpu_address_t *src)
 {
+    uvm_pushbuffer_t *pushbuffer;
+
     if (!uvm_channel_is_proxy(push->channel))
         return;
 
-    src->address -= uvm_pushbuffer_get_gpu_va_for_push(push->channel->pool->manager->pushbuffer, push);
+    pushbuffer = uvm_channel_get_pushbuffer(push->channel);
+    src->address -= uvm_pushbuffer_get_gpu_va_for_push(pushbuffer, push);
 }
 
 bool uvm_hal_ampere_ce_memset_is_valid_c6b5(uvm_push_t *push,

kernel-open/nvidia-uvm/uvm_ats.c

@@ -44,6 +44,8 @@ void uvm_ats_init(const UvmPlatformInfo *platform_info)
 
 void uvm_ats_init_va_space(uvm_va_space_t *va_space)
 {
+    uvm_init_rwsem(&va_space->ats.lock, UVM_LOCK_ORDER_LEAF);
+
     if (UVM_ATS_IBM_SUPPORTED())
         uvm_ats_ibm_init_va_space(va_space);
 }

kernel-open/nvidia-uvm/uvm_ats.h

@@ -28,6 +28,7 @@
 #include "uvm_forward_decl.h"
 #include "uvm_ats_ibm.h"
 #include "nv_uvm_types.h"
+#include "uvm_lock.h"
 
     #include "uvm_ats_sva.h"
 
@@ -39,6 +40,10 @@ typedef struct
     // indexed by gpu->id. This mask is protected by the VA space lock.
     uvm_processor_mask_t registered_gpu_va_spaces;
 
+    // Protects racing invalidates in the VA space while hmm_range_fault() is
+    // being called in ats_compute_residency_mask().
+    uvm_rw_semaphore_t lock;
+
     union
     {
         uvm_ibm_va_space_t ibm;

kernel-open/nvidia-uvm/uvm_ats_faults.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2018 NVIDIA Corporation
+    Copyright (c) 2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -20,74 +20,46 @@
     DEALINGS IN THE SOFTWARE.
 *******************************************************************************/
 
+#include "uvm_api.h"
 #include "uvm_tools.h"
 #include "uvm_va_range.h"
+#include "uvm_ats.h"
 #include "uvm_ats_faults.h"
 #include "uvm_migrate_pageable.h"
+#include <linux/nodemask.h>
 #include <linux/mempolicy.h>
+#include <linux/mmu_notifier.h>
 
-// TODO: Bug 2103669: Implement a real prefetching policy and remove or adapt
-// these experimental parameters. These are intended to help guide that policy.
-static unsigned int uvm_exp_perf_prefetch_ats_order_replayable = 0;
-module_param(uvm_exp_perf_prefetch_ats_order_replayable, uint, 0644);
-MODULE_PARM_DESC(uvm_exp_perf_prefetch_ats_order_replayable,
-                 "Max order of pages (2^N) to prefetch on replayable ATS faults");
+#if UVM_HMM_RANGE_FAULT_SUPPORTED()
+#include <linux/hmm.h>
+#endif
 
-static unsigned int uvm_exp_perf_prefetch_ats_order_non_replayable = 0;
-module_param(uvm_exp_perf_prefetch_ats_order_non_replayable, uint, 0644);
-MODULE_PARM_DESC(uvm_exp_perf_prefetch_ats_order_non_replayable,
-                 "Max order of pages (2^N) to prefetch on non-replayable ATS faults");
-
-// Expand the fault region to the naturally-aligned region with order given by
-// the module parameters, clamped to the vma containing fault_addr (if any).
-// Note that this means the region contains fault_addr but may not begin at
-// fault_addr.
-static void expand_fault_region(struct vm_area_struct *vma,
-                                NvU64 start,
-                                size_t length,
-                                uvm_fault_client_type_t client_type,
-                                unsigned long *migrate_start,
-                                unsigned long *migrate_length)
+typedef enum
 {
-    unsigned int order;
-    unsigned long outer, aligned_start, aligned_size;
+    UVM_ATS_SERVICE_TYPE_FAULTS = 0,
+    UVM_ATS_SERVICE_TYPE_ACCESS_COUNTERS,
+    UVM_ATS_SERVICE_TYPE_COUNT
+} uvm_ats_service_type_t;
 
-    *migrate_start = start;
-    *migrate_length = length;
-
-    if (client_type == UVM_FAULT_CLIENT_TYPE_HUB)
-        order = uvm_exp_perf_prefetch_ats_order_non_replayable;
-    else
-        order = uvm_exp_perf_prefetch_ats_order_replayable;
-
-    if (order == 0)
-        return;
-
-    UVM_ASSERT(vma);
-    UVM_ASSERT(order < BITS_PER_LONG - PAGE_SHIFT);
-
-    aligned_size = (1UL << order) * PAGE_SIZE;
-
-    aligned_start = start & ~(aligned_size - 1);
-
-    *migrate_start = max(vma->vm_start, aligned_start);
-    outer = min(vma->vm_end, aligned_start + aligned_size);
-    *migrate_length = outer - *migrate_start;
-}
-
-static NV_STATUS service_ats_faults(uvm_gpu_va_space_t *gpu_va_space,
-                                    struct vm_area_struct *vma,
-                                    NvU64 start,
-                                    size_t length,
-                                    uvm_fault_access_type_t access_type,
-                                    uvm_ats_fault_context_t *ats_context)
+static NV_STATUS service_ats_requests(uvm_gpu_va_space_t *gpu_va_space,
+                                      struct vm_area_struct *vma,
+                                      NvU64 start,
+                                      size_t length,
+                                      uvm_fault_access_type_t access_type,
+                                      uvm_ats_service_type_t service_type,
+                                      uvm_ats_fault_context_t *ats_context)
 {
     uvm_va_space_t *va_space = gpu_va_space->va_space;
     struct mm_struct *mm = va_space->va_space_mm.mm;
-    bool write = (access_type >= UVM_FAULT_ACCESS_TYPE_WRITE);
     NV_STATUS status;
     NvU64 user_space_start;
     NvU64 user_space_length;
+    bool write = (access_type >= UVM_FAULT_ACCESS_TYPE_WRITE);
+    bool fault_service_type = (service_type == UVM_ATS_SERVICE_TYPE_FAULTS);
+    uvm_populate_permissions_t populate_permissions = fault_service_type ?
+                                            (write ? UVM_POPULATE_PERMISSIONS_WRITE : UVM_POPULATE_PERMISSIONS_ANY) :
+                                            UVM_POPULATE_PERMISSIONS_INHERIT;
+
 
     // Request uvm_migrate_pageable() to touch the corresponding page after
     // population.
@@ -122,23 +94,18 @@ static NV_STATUS service_ats_faults(uvm_gpu_va_space_t *gpu_va_space,
         .mm                             = mm,
         .dst_id                         = ats_context->residency_id,
         .dst_node_id                    = ats_context->residency_node,
-        .populate_permissions           = write ? UVM_POPULATE_PERMISSIONS_WRITE : UVM_POPULATE_PERMISSIONS_ANY,
-        .touch                          = true,
-        .skip_mapped                    = true,
-        .populate_on_cpu_alloc_failures = true,
+        .start                          = start,
+        .length                         = length,
+        .populate_permissions           = populate_permissions,
+        .touch                          = fault_service_type,
+        .skip_mapped                    = fault_service_type,
+        .populate_on_cpu_alloc_failures = fault_service_type,
         .user_space_start               = &user_space_start,
         .user_space_length              = &user_space_length,
     };
 
     UVM_ASSERT(uvm_ats_can_service_faults(gpu_va_space, mm));
 
-    expand_fault_region(vma,
-                        start,
-                        length,
-                        ats_context->client_type,
-                        &uvm_migrate_args.start,
-                        &uvm_migrate_args.length);
-
     // We are trying to use migrate_vma API in the kernel (if it exists) to
     // populate and map the faulting region on the GPU. We want to do this only
     // on the first touch. That is, pages which are not already mapped. So, we
@@ -153,10 +120,10 @@ static NV_STATUS service_ats_faults(uvm_gpu_va_space_t *gpu_va_space,
     return status;
 }
 
-static void flush_tlb_write_faults(uvm_gpu_va_space_t *gpu_va_space,
-                                   NvU64 addr,
-                                   size_t size,
-                                   uvm_fault_client_type_t client_type)
+static void flush_tlb_va_region(uvm_gpu_va_space_t *gpu_va_space,
+                                NvU64 addr,
+                                size_t size,
+                                uvm_fault_client_type_t client_type)
 {
     uvm_ats_fault_invalidate_t *ats_invalidate;
 
@@ -165,12 +132,12 @@ static void flush_tlb_write_faults(uvm_gpu_va_space_t *gpu_va_space,
     else
         ats_invalidate = &gpu_va_space->gpu->parent->fault_buffer_info.non_replayable.ats_invalidate;
 
-    if (!ats_invalidate->write_faults_in_batch) {
-        uvm_tlb_batch_begin(&gpu_va_space->page_tables, &ats_invalidate->write_faults_tlb_batch);
-        ats_invalidate->write_faults_in_batch = true;
+    if (!ats_invalidate->tlb_batch_pending) {
+        uvm_tlb_batch_begin(&gpu_va_space->page_tables, &ats_invalidate->tlb_batch);
+        ats_invalidate->tlb_batch_pending = true;
     }
 
-    uvm_tlb_batch_invalidate(&ats_invalidate->write_faults_tlb_batch, addr, size, PAGE_SIZE, UVM_MEMBAR_NONE);
+    uvm_tlb_batch_invalidate(&ats_invalidate->tlb_batch, addr, size, PAGE_SIZE, UVM_MEMBAR_NONE);
 }
 
 static void ats_batch_select_residency(uvm_gpu_va_space_t *gpu_va_space,
@@ -184,12 +151,22 @@ static void ats_batch_select_residency(uvm_gpu_va_space_t *gpu_va_space,
     struct mempolicy *vma_policy = vma_policy(vma);
     unsigned short mode;
 
+    ats_context->prefetch_state.has_preferred_location = false;
+
+    // It's safe to read vma_policy since the mmap_lock is held in at least read
+    // mode in this path.
+    uvm_assert_mmap_lock_locked(vma->vm_mm);
+
     if (!vma_policy)
         goto done;
 
     mode = vma_policy->mode;
 
-    if ((mode == MPOL_BIND) || (mode == MPOL_PREFERRED_MANY) || (mode == MPOL_PREFERRED)) {
+    if ((mode == MPOL_BIND)
+#if defined(NV_MPOL_PREFERRED_MANY_PRESENT)
+         || (mode == MPOL_PREFERRED_MANY)
+#endif
+         || (mode == MPOL_PREFERRED)) {
         int home_node = NUMA_NO_NODE;
 
 #if defined(NV_MEMPOLICY_HAS_HOME_NODE)
@@ -212,6 +189,9 @@ static void ats_batch_select_residency(uvm_gpu_va_space_t *gpu_va_space,
             else
                 residency = first_node(vma_policy->nodes);
         }
+
+        if (!nodes_empty(vma_policy->nodes))
+            ats_context->prefetch_state.has_preferred_location = true;
     }
 
     // Update gpu if residency is not the faulting gpu.
@@ -219,12 +199,258 @@ static void ats_batch_select_residency(uvm_gpu_va_space_t *gpu_va_space,
         gpu = uvm_va_space_find_gpu_with_memory_node_id(gpu_va_space->va_space, residency);
 
 done:
+#else
+    ats_context->prefetch_state.has_preferred_location = false;
 #endif
 
     ats_context->residency_id = gpu ? gpu->parent->id : UVM_ID_CPU;
     ats_context->residency_node = residency;
 }
 
+static void get_range_in_vma(struct vm_area_struct *vma, NvU64 base, NvU64 *start, NvU64 *end)
+{
+    *start = max(vma->vm_start, (unsigned long) base);
+    *end = min(vma->vm_end, (unsigned long) (base + UVM_VA_BLOCK_SIZE));
+}
+
+static uvm_page_index_t uvm_ats_cpu_page_index(NvU64 base, NvU64 addr)
+{
+    UVM_ASSERT(addr >= base);
+    UVM_ASSERT(addr <= (base + UVM_VA_BLOCK_SIZE));
+
+    return (addr - base) / PAGE_SIZE;
+}
+
+// start and end must be aligned to PAGE_SIZE and must fall within
+// [base, base + UVM_VA_BLOCK_SIZE]
+static uvm_va_block_region_t uvm_ats_region_from_start_end(NvU64 start, NvU64 end)
+{
+    // base can be greater than, less than or equal to the start of a VMA.
+    NvU64 base = UVM_VA_BLOCK_ALIGN_DOWN(start);
+
+    UVM_ASSERT(start < end);
+    UVM_ASSERT(PAGE_ALIGNED(start));
+    UVM_ASSERT(PAGE_ALIGNED(end));
+    UVM_ASSERT(IS_ALIGNED(base, UVM_VA_BLOCK_SIZE));
+
+    return uvm_va_block_region(uvm_ats_cpu_page_index(base, start), uvm_ats_cpu_page_index(base, end));
+}
+
+static uvm_va_block_region_t uvm_ats_region_from_vma(struct vm_area_struct *vma, NvU64 base)
+{
+    NvU64 start;
+    NvU64 end;
+
+    get_range_in_vma(vma, base, &start, &end);
+
+    return uvm_ats_region_from_start_end(start, end);
+}
+
+#if UVM_HMM_RANGE_FAULT_SUPPORTED()
+
+static bool uvm_ats_invalidate_notifier(struct mmu_interval_notifier *mni, unsigned long cur_seq)
+{
+    uvm_ats_fault_context_t *ats_context = container_of(mni, uvm_ats_fault_context_t, prefetch_state.notifier);
+    uvm_va_space_t *va_space = ats_context->prefetch_state.va_space;
+
+    // The following write lock protects against concurrent invalidates while
+    // hmm_range_fault() is being called in ats_compute_residency_mask().
+    uvm_down_write(&va_space->ats.lock);
+
+    mmu_interval_set_seq(mni, cur_seq);
+
+    uvm_up_write(&va_space->ats.lock);
+
+    return true;
+}
+
+static bool uvm_ats_invalidate_notifier_entry(struct mmu_interval_notifier *mni,
+                                              const struct mmu_notifier_range *range,
+                                              unsigned long cur_seq)
+{
+    UVM_ENTRY_RET(uvm_ats_invalidate_notifier(mni, cur_seq));
+}
+
+static const struct mmu_interval_notifier_ops uvm_ats_notifier_ops =
+{
+    .invalidate = uvm_ats_invalidate_notifier_entry,
+};
+
+#endif
+
+static NV_STATUS ats_compute_residency_mask(uvm_gpu_va_space_t *gpu_va_space,
+                                            struct vm_area_struct *vma,
+                                            NvU64 base,
+                                            uvm_ats_fault_context_t *ats_context)
+{
+    NV_STATUS status = NV_OK;
+    uvm_page_mask_t *residency_mask = &ats_context->prefetch_state.residency_mask;
+
+#if UVM_HMM_RANGE_FAULT_SUPPORTED()
+    int ret;
+    NvU64 start;
+    NvU64 end;
+    struct hmm_range range;
+    uvm_page_index_t page_index;
+    uvm_va_block_region_t vma_region;
+    uvm_va_space_t *va_space = gpu_va_space->va_space;
+    struct mm_struct *mm = va_space->va_space_mm.mm;
+
+    uvm_assert_rwsem_locked_read(&va_space->lock);
+
+    ats_context->prefetch_state.first_touch = true;
+
+    uvm_page_mask_zero(residency_mask);
+
+    get_range_in_vma(vma, base, &start, &end);
+
+    vma_region = uvm_ats_region_from_start_end(start, end);
+
+    range.notifier = &ats_context->prefetch_state.notifier;
+    range.start = start;
+    range.end = end;
+    range.hmm_pfns = ats_context->prefetch_state.pfns;
+    range.default_flags = 0;
+    range.pfn_flags_mask = 0;
+    range.dev_private_owner = NULL;
+
+    ats_context->prefetch_state.va_space = va_space;
+
+    // mmu_interval_notifier_insert() will try to acquire mmap_lock for write
+    // and will deadlock since mmap_lock is already held for read in this path.
+    // This is prevented by calling __mmu_notifier_register() during va_space
+    // creation. See the comment in uvm_mmu_notifier_register() for more
+    // details.
+    ret = mmu_interval_notifier_insert(range.notifier, mm, start, end, &uvm_ats_notifier_ops);
+    if (ret)
+        return errno_to_nv_status(ret);
+
+    while (true) {
+        range.notifier_seq = mmu_interval_read_begin(range.notifier);
+        ret = hmm_range_fault(&range);
+        if (ret == -EBUSY)
+            continue;
+        if (ret) {
+            status = errno_to_nv_status(ret);
+            UVM_ASSERT(status != NV_OK);
+            break;
+        }
+
+        uvm_down_read(&va_space->ats.lock);
+
+        // Pages may have been freed or re-allocated after hmm_range_fault() is
+        // called. So the PTE might point to a different page or nothing. In the
+        // memory hot-unplug case it is not safe to call page_to_nid() on the
+        // page as the struct page itself may have been freed. To protect
+        // against these cases, uvm_ats_invalidate_entry() blocks on va_space
+        // ATS write lock for concurrent invalidates since va_space ATS lock is
+        // held for read in this path.
+        if (!mmu_interval_read_retry(range.notifier, range.notifier_seq))
+            break;
+
+        uvm_up_read(&va_space->ats.lock);
+    }
+
+    if (status == NV_OK) {
+        for_each_va_block_page_in_region(page_index, vma_region) {
+            unsigned long pfn = ats_context->prefetch_state.pfns[page_index - vma_region.first];
+
+            if (pfn & HMM_PFN_VALID) {
+                struct page *page = hmm_pfn_to_page(pfn);
+
+                if (page_to_nid(page) == ats_context->residency_node)
+                    uvm_page_mask_set(residency_mask, page_index);
+
+                ats_context->prefetch_state.first_touch = false;
+            }
+        }
+
+        uvm_up_read(&va_space->ats.lock);
+    }
+
+    mmu_interval_notifier_remove(range.notifier);
+
+#else
+    uvm_page_mask_zero(residency_mask);
+#endif
+
+    return status;
+}
+
+static void ats_compute_prefetch_mask(uvm_gpu_va_space_t *gpu_va_space,
+                                      struct vm_area_struct *vma,
+                                      uvm_ats_fault_context_t *ats_context,
+                                      uvm_va_block_region_t max_prefetch_region)
+{
+    uvm_page_mask_t *accessed_mask = &ats_context->accessed_mask;
+    uvm_page_mask_t *residency_mask = &ats_context->prefetch_state.residency_mask;
+    uvm_page_mask_t *prefetch_mask = &ats_context->prefetch_state.prefetch_pages_mask;
+    uvm_perf_prefetch_bitmap_tree_t *bitmap_tree = &ats_context->prefetch_state.bitmap_tree;
+
+    if (uvm_page_mask_empty(accessed_mask))
+        return;
+
+    uvm_perf_prefetch_compute_ats(gpu_va_space->va_space,
+                                  accessed_mask,
+                                  uvm_va_block_region_from_mask(NULL, accessed_mask),
+                                  max_prefetch_region,
+                                  residency_mask,
+                                  bitmap_tree,
+                                  prefetch_mask);
+}
+
+static NV_STATUS ats_compute_prefetch(uvm_gpu_va_space_t *gpu_va_space,
+                                      struct vm_area_struct *vma,
+                                      NvU64 base,
+                                      uvm_ats_service_type_t service_type,
+                                      uvm_ats_fault_context_t *ats_context)
+{
+    NV_STATUS status;
+    uvm_page_mask_t *accessed_mask = &ats_context->accessed_mask;
+    uvm_page_mask_t *prefetch_mask = &ats_context->prefetch_state.prefetch_pages_mask;
+    uvm_va_block_region_t max_prefetch_region = uvm_ats_region_from_vma(vma, base);
+
+    // Residency mask needs to be computed even if prefetching is disabled since
+    // the residency information is also needed by access counters servicing in
+    // uvm_ats_service_access_counters()
+    status = ats_compute_residency_mask(gpu_va_space, vma, base, ats_context);
+    if (status != NV_OK)
+        return status;
+
+    if (!uvm_perf_prefetch_enabled(gpu_va_space->va_space))
+        return status;
+
+    if (uvm_page_mask_empty(accessed_mask))
+        return status;
+
+    // Prefetch the entire region if none of the pages are resident on any node
+    // and if preferred_location is the faulting GPU.
+    if (ats_context->prefetch_state.has_preferred_location &&
+        (ats_context->prefetch_state.first_touch || (service_type == UVM_ATS_SERVICE_TYPE_ACCESS_COUNTERS)) &&
+        uvm_id_equal(ats_context->residency_id, gpu_va_space->gpu->id)) {
+
+        uvm_page_mask_init_from_region(prefetch_mask, max_prefetch_region, NULL);
+    }
+    else {
+        ats_compute_prefetch_mask(gpu_va_space, vma, ats_context, max_prefetch_region);
+    }
+
+    if (service_type == UVM_ATS_SERVICE_TYPE_FAULTS) {
+        uvm_page_mask_t *read_fault_mask = &ats_context->read_fault_mask;
+        uvm_page_mask_t *write_fault_mask = &ats_context->write_fault_mask;
+
+        uvm_page_mask_or(read_fault_mask, read_fault_mask, prefetch_mask);
+
+        if (vma->vm_flags & VM_WRITE)
+            uvm_page_mask_or(write_fault_mask, write_fault_mask, prefetch_mask);
+    }
+    else {
+        uvm_page_mask_or(accessed_mask, accessed_mask, prefetch_mask);
+    }
+
+    return status;
+}
+
 NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
                                  struct vm_area_struct *vma,
                                  NvU64 base,
@@ -238,6 +464,7 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
     uvm_page_mask_t *faults_serviced_mask = &ats_context->faults_serviced_mask;
     uvm_page_mask_t *reads_serviced_mask = &ats_context->reads_serviced_mask;
     uvm_fault_client_type_t client_type = ats_context->client_type;
+    uvm_ats_service_type_t service_type = UVM_ATS_SERVICE_TYPE_FAULTS;
 
     UVM_ASSERT(vma);
     UVM_ASSERT(IS_ALIGNED(base, UVM_VA_BLOCK_SIZE));
@@ -246,6 +473,9 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
     UVM_ASSERT(gpu_va_space->ats.enabled);
     UVM_ASSERT(uvm_gpu_va_space_state(gpu_va_space) == UVM_GPU_VA_SPACE_STATE_ACTIVE);
 
+    uvm_assert_mmap_lock_locked(vma->vm_mm);
+    uvm_assert_rwsem_locked(&gpu_va_space->va_space->lock);
+
     uvm_page_mask_zero(faults_serviced_mask);
     uvm_page_mask_zero(reads_serviced_mask);
 
@@ -263,10 +493,16 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
             uvm_page_mask_and(write_fault_mask, write_fault_mask, read_fault_mask);
         else
             uvm_page_mask_zero(write_fault_mask);
+
+        // There are no pending faults beyond write faults to RO region.
+        if (uvm_page_mask_empty(read_fault_mask))
+            return status;
     }
 
     ats_batch_select_residency(gpu_va_space, vma, ats_context);
 
+    ats_compute_prefetch(gpu_va_space, vma, base, service_type, ats_context);
+
     for_each_va_block_subregion_in_mask(subregion, write_fault_mask, region) {
         NvU64 start = base + (subregion.first * PAGE_SIZE);
         size_t length = uvm_va_block_region_num_pages(subregion) * PAGE_SIZE;
@@ -277,12 +513,13 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
         UVM_ASSERT(start >= vma->vm_start);
         UVM_ASSERT((start + length) <= vma->vm_end);
 
-        status = service_ats_faults(gpu_va_space, vma, start, length, access_type, ats_context);
+        status = service_ats_requests(gpu_va_space, vma, start, length, access_type, service_type, ats_context);
         if (status != NV_OK)
             return status;
 
         if (vma->vm_flags & VM_WRITE) {
             uvm_page_mask_region_fill(faults_serviced_mask, subregion);
+            uvm_ats_smmu_invalidate_tlbs(gpu_va_space, start, length);
 
             // The Linux kernel never invalidates TLB entries on mapping
             // permission upgrade. This is a problem if the GPU has cached
@@ -293,7 +530,7 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
             // infinite loop because we just forward the fault to the Linux
             // kernel and it will see that the permissions in the page table are
             // correct. Therefore, we flush TLB entries on ATS write faults.
-            flush_tlb_write_faults(gpu_va_space, start, length, client_type);
+            flush_tlb_va_region(gpu_va_space, start, length, client_type);
         }
         else {
             uvm_page_mask_region_fill(reads_serviced_mask, subregion);
@@ -311,11 +548,20 @@ NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
         UVM_ASSERT(start >= vma->vm_start);
         UVM_ASSERT((start + length) <= vma->vm_end);
 
-        status = service_ats_faults(gpu_va_space, vma, start, length, access_type, ats_context);
+        status = service_ats_requests(gpu_va_space, vma, start, length, access_type, service_type, ats_context);
         if (status != NV_OK)
             return status;
 
         uvm_page_mask_region_fill(faults_serviced_mask, subregion);
+
+        // Similarly to permission upgrade scenario, discussed above, GPU
+        // will not re-fetch the entry if the PTE is invalid and page size
+        // is 4K. To avoid infinite faulting loop, invalidate TLB for every
+        // new translation written explicitly like in the case of permission
+        // upgrade.
+        if (PAGE_SIZE == UVM_PAGE_SIZE_4K)
+            flush_tlb_va_region(gpu_va_space, start, length, client_type);
+
     }
 
     return status;
@@ -350,7 +596,7 @@ NV_STATUS uvm_ats_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space,
     NV_STATUS status;
     uvm_push_t push;
 
-    if (!ats_invalidate->write_faults_in_batch)
+    if (!ats_invalidate->tlb_batch_pending)
         return NV_OK;
 
     UVM_ASSERT(gpu_va_space);
@@ -362,7 +608,7 @@ NV_STATUS uvm_ats_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space,
                             "Invalidate ATS entries");
 
     if (status == NV_OK) {
-        uvm_tlb_batch_end(&ats_invalidate->write_faults_tlb_batch, &push, UVM_MEMBAR_NONE);
+        uvm_tlb_batch_end(&ats_invalidate->tlb_batch, &push, UVM_MEMBAR_NONE);
         uvm_push_end(&push);
 
         // Add this push to the GPU's tracker so that fault replays/clears can
@@ -370,8 +616,57 @@ NV_STATUS uvm_ats_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space,
         status = uvm_tracker_add_push_safe(out_tracker, &push);
     }
 
-    ats_invalidate->write_faults_in_batch = false;
+    ats_invalidate->tlb_batch_pending = false;
 
     return status;
 }
 
+NV_STATUS uvm_ats_service_access_counters(uvm_gpu_va_space_t *gpu_va_space,
+                                          struct vm_area_struct *vma,
+                                          NvU64 base,
+                                          uvm_ats_fault_context_t *ats_context)
+{
+    uvm_va_block_region_t subregion;
+    uvm_va_block_region_t region = uvm_va_block_region(0, PAGES_PER_UVM_VA_BLOCK);
+    uvm_ats_service_type_t service_type = UVM_ATS_SERVICE_TYPE_ACCESS_COUNTERS;
+
+    UVM_ASSERT(vma);
+    UVM_ASSERT(IS_ALIGNED(base, UVM_VA_BLOCK_SIZE));
+    UVM_ASSERT(g_uvm_global.ats.enabled);
+    UVM_ASSERT(gpu_va_space);
+    UVM_ASSERT(gpu_va_space->ats.enabled);
+    UVM_ASSERT(uvm_gpu_va_space_state(gpu_va_space) == UVM_GPU_VA_SPACE_STATE_ACTIVE);
+
+    uvm_assert_mmap_lock_locked(vma->vm_mm);
+    uvm_assert_rwsem_locked(&gpu_va_space->va_space->lock);
+
+    ats_batch_select_residency(gpu_va_space, vma, ats_context);
+
+    // Ignoring the return value of ats_compute_prefetch is ok since prefetching
+    // is just an optimization and servicing access counter migrations is still
+    // worthwhile even without any prefetching added. So, let servicing continue
+    // instead of returning early even if the prefetch computation fails.
+    ats_compute_prefetch(gpu_va_space, vma, base, service_type, ats_context);
+
+    // Remove pages which are already resident at the intended destination from
+    // the accessed_mask.
+    uvm_page_mask_andnot(&ats_context->accessed_mask,
+                         &ats_context->accessed_mask,
+                         &ats_context->prefetch_state.residency_mask);
+
+    for_each_va_block_subregion_in_mask(subregion, &ats_context->accessed_mask, region) {
+        NV_STATUS status;
+        NvU64 start = base + (subregion.first * PAGE_SIZE);
+        size_t length = uvm_va_block_region_num_pages(subregion) * PAGE_SIZE;
+        uvm_fault_access_type_t access_type = UVM_FAULT_ACCESS_TYPE_COUNT;
+
+        UVM_ASSERT(start >= vma->vm_start);
+        UVM_ASSERT((start + length) <= vma->vm_end);
+
+        status = service_ats_requests(gpu_va_space, vma, start, length, access_type, service_type, ats_context);
+        if (status != NV_OK)
+            return status;
+    }
+
+    return NV_OK;
+}

kernel-open/nvidia-uvm/uvm_ats_faults.h

@@ -42,17 +42,37 @@
 // corresponding bit in read_fault_mask. These returned masks are only valid if
 // the return status is NV_OK. Status other than NV_OK indicate system global
 // fault servicing failures.
+//
+// LOCKING: The caller must retain and hold the mmap_lock and hold the va_space
+// lock.
 NV_STATUS uvm_ats_service_faults(uvm_gpu_va_space_t *gpu_va_space,
                                  struct vm_area_struct *vma,
                                  NvU64 base,
                                  uvm_ats_fault_context_t *ats_context);
 
+// Service access counter notifications on ATS regions in the range (base, base
+// + UVM_VA_BLOCK_SIZE) for individual pages in the range requested by page_mask
+// set in ats_context->accessed_mask. base must be aligned to UVM_VA_BLOCK_SIZE.
+// The caller is responsible for ensuring that the addresses in the
+// accessed_mask is completely covered by the VMA. The caller is also
+// responsible for handling any errors returned by this function.
+//
+// Returns NV_OK if servicing was successful. Any other error indicates an error
+// while servicing the range.
+//
+// LOCKING: The caller must retain and hold the mmap_lock and hold the va_space
+// lock.
+NV_STATUS uvm_ats_service_access_counters(uvm_gpu_va_space_t *gpu_va_space,
+                                          struct vm_area_struct *vma,
+                                          NvU64 base,
+                                          uvm_ats_fault_context_t *ats_context);
+
 // Return whether there are any VA ranges (and thus GMMU mappings) within the
 // UVM_GMMU_ATS_GRANULARITY-aligned region containing address.
 bool uvm_ats_check_in_gmmu_region(uvm_va_space_t *va_space, NvU64 address, uvm_va_range_t *next);
 
 // This function performs pending TLB invalidations for ATS and clears the
-// ats_invalidate->write_faults_in_batch flag
+// ats_invalidate->tlb_batch_pending flag
 NV_STATUS uvm_ats_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space,
                                   uvm_ats_fault_invalidate_t *ats_invalidate,
                                   uvm_tracker_t *out_tracker);

kernel-open/nvidia-uvm/uvm_ats_sva.c

@@ -29,8 +29,13 @@
 #include "uvm_va_space.h"
 #include "uvm_va_space_mm.h"
 
+#include <asm/io.h>
+#include <linux/log2.h>
 #include <linux/iommu.h>
 #include <linux/mm_types.h>
+#include <linux/acpi.h>
+#include <linux/device.h>
+#include <linux/mmu_context.h>
 
 // linux/sched/mm.h is needed for mmget_not_zero and mmput to get the mm
 // reference required for the iommu_sva_bind_device() call. This header is not
@@ -46,18 +51,284 @@
 #define UVM_IOMMU_SVA_BIND_DEVICE(dev, mm) iommu_sva_bind_device(dev, mm)
 #endif
 
+// Type to represent a 128-bit SMMU command queue command.
+struct smmu_cmd {
+    NvU64 low;
+    NvU64 high;
+};
+
+// Base address of SMMU CMDQ-V for GSMMU0.
+#define SMMU_CMDQV_BASE_ADDR(smmu_base) (smmu_base + 0x200000)
+#define SMMU_CMDQV_BASE_LEN 0x00830000
+
+// CMDQV configuration is done by firmware but we check status here.
+#define SMMU_CMDQV_CONFIG 0x0
+#define SMMU_CMDQV_CONFIG_CMDQV_EN BIT(0)
+
+// Used to map a particular VCMDQ to a VINTF.
+#define SMMU_CMDQV_CMDQ_ALLOC_MAP(vcmdq_id) (0x200 + 0x4 * (vcmdq_id))
+#define SMMU_CMDQV_CMDQ_ALLOC_MAP_ALLOC BIT(0)
+
+// Shift for the field containing the index of the virtual interface
+// owning the VCMDQ.
+#define SMMU_CMDQV_CMDQ_ALLOC_MAP_VIRT_INTF_INDX_SHIFT 15
+
+// Base address for the VINTF registers.
+#define SMMU_VINTF_BASE_ADDR(cmdqv_base_addr, vintf_id) (cmdqv_base_addr + 0x1000 + 0x100 * (vintf_id))
+
+// Virtual interface (VINTF) configuration registers. The WAR only
+// works on baremetal so we need to configure ourselves as the
+// hypervisor owner.
+#define SMMU_VINTF_CONFIG 0x0
+#define SMMU_VINTF_CONFIG_ENABLE BIT(0)
+#define SMMU_VINTF_CONFIG_HYP_OWN BIT(17)
+
+#define SMMU_VINTF_STATUS 0x0
+#define SMMU_VINTF_STATUS_ENABLED BIT(0)
+
+// Caclulates the base address for a particular VCMDQ instance.
+#define SMMU_VCMDQ_BASE_ADDR(cmdqv_base_addr, vcmdq_id) (cmdqv_base_addr + 0x10000 + 0x80 * (vcmdq_id))
+
+// SMMU command queue consumer index register. Updated by SMMU
+// when commands are consumed.
+#define SMMU_VCMDQ_CONS 0x0
+
+// SMMU command queue producer index register. Updated by UVM when
+// commands are added to the queue.
+#define SMMU_VCMDQ_PROD 0x4
+
+// Configuration register used to enable a VCMDQ.
+#define SMMU_VCMDQ_CONFIG 0x8
+#define SMMU_VCMDQ_CONFIG_ENABLE BIT(0)
+
+// Status register used to check the VCMDQ is enabled.
+#define SMMU_VCMDQ_STATUS 0xc
+#define SMMU_VCMDQ_STATUS_ENABLED BIT(0)
+
+// Base address offset for the VCMDQ registers.
+#define SMMU_VCMDQ_CMDQ_BASE 0x10000
+
+// Size of the command queue. Each command is 16 bytes and we can't
+// have a command queue greater than one page in size.
+#define SMMU_VCMDQ_CMDQ_BASE_LOG2SIZE (PAGE_SHIFT - ilog2(sizeof(struct smmu_cmd)))
+#define SMMU_VCMDQ_CMDQ_ENTRIES (1UL << SMMU_VCMDQ_CMDQ_BASE_LOG2SIZE)
+
+// We always use VINTF63 for the WAR
+#define VINTF 63
+static void smmu_vintf_write32(void __iomem *smmu_cmdqv_base, int reg, NvU32 val)
+{
+    iowrite32(val, SMMU_VINTF_BASE_ADDR(smmu_cmdqv_base, VINTF) + reg);
+}
+
+static NvU32 smmu_vintf_read32(void __iomem *smmu_cmdqv_base, int reg)
+{
+    return ioread32(SMMU_VINTF_BASE_ADDR(smmu_cmdqv_base, VINTF) + reg);
+}
+
+// We always use VCMDQ127 for the WAR
+#define VCMDQ 127
+void smmu_vcmdq_write32(void __iomem *smmu_cmdqv_base, int reg, NvU32 val)
+{
+    iowrite32(val, SMMU_VCMDQ_BASE_ADDR(smmu_cmdqv_base, VCMDQ) + reg);
+}
+
+NvU32 smmu_vcmdq_read32(void __iomem *smmu_cmdqv_base, int reg)
+{
+    return ioread32(SMMU_VCMDQ_BASE_ADDR(smmu_cmdqv_base, VCMDQ) + reg);
+}
+
+static void smmu_vcmdq_write64(void __iomem *smmu_cmdqv_base, int reg, NvU64 val)
+{
+#if NV_IS_EXPORT_SYMBOL_PRESENT___iowrite64_lo_hi
+    __iowrite64_lo_hi(val, SMMU_VCMDQ_BASE_ADDR(smmu_cmdqv_base, VCMDQ) + reg);
+#else
+    iowrite64(val, SMMU_VCMDQ_BASE_ADDR(smmu_cmdqv_base, VCMDQ) + reg);
+#endif
+}
+
+// Fix for Bug 4130089: [GH180][r535] WAR for kernel not issuing SMMU
+// TLB invalidates on read-only to read-write upgrades
+static NV_STATUS uvm_ats_smmu_war_init(uvm_parent_gpu_t *parent_gpu)
+{
+    uvm_spin_loop_t spin;
+    NV_STATUS status;
+    unsigned long cmdqv_config;
+    void __iomem *smmu_cmdqv_base;
+    struct acpi_iort_node *node;
+    struct acpi_iort_smmu_v3 *iort_smmu;
+
+    node = *(struct acpi_iort_node **) dev_get_platdata(parent_gpu->pci_dev->dev.iommu->iommu_dev->dev->parent);
+    iort_smmu = (struct acpi_iort_smmu_v3 *) node->node_data;
+
+    smmu_cmdqv_base = ioremap(SMMU_CMDQV_BASE_ADDR(iort_smmu->base_address), SMMU_CMDQV_BASE_LEN);
+    if (!smmu_cmdqv_base)
+        return NV_ERR_NO_MEMORY;
+
+    parent_gpu->smmu_war.smmu_cmdqv_base = smmu_cmdqv_base;
+    cmdqv_config = ioread32(smmu_cmdqv_base + SMMU_CMDQV_CONFIG);
+    if (!(cmdqv_config & SMMU_CMDQV_CONFIG_CMDQV_EN)) {
+        status = NV_ERR_OBJECT_NOT_FOUND;
+        goto out;
+    }
+
+    // Allocate SMMU CMDQ pages for WAR
+    parent_gpu->smmu_war.smmu_cmdq = alloc_page(NV_UVM_GFP_FLAGS | __GFP_ZERO);
+    if (!parent_gpu->smmu_war.smmu_cmdq) {
+        status = NV_ERR_NO_MEMORY;
+        goto out;
+    }
+
+    // Initialise VINTF for the WAR
+    smmu_vintf_write32(smmu_cmdqv_base, SMMU_VINTF_CONFIG, SMMU_VINTF_CONFIG_ENABLE | SMMU_VINTF_CONFIG_HYP_OWN);
+    UVM_SPIN_WHILE(!(smmu_vintf_read32(smmu_cmdqv_base, SMMU_VINTF_STATUS) & SMMU_VINTF_STATUS_ENABLED), &spin);
+
+    // Allocate VCMDQ to VINTF
+    iowrite32((VINTF << SMMU_CMDQV_CMDQ_ALLOC_MAP_VIRT_INTF_INDX_SHIFT) | SMMU_CMDQV_CMDQ_ALLOC_MAP_ALLOC,
+              smmu_cmdqv_base + SMMU_CMDQV_CMDQ_ALLOC_MAP(VCMDQ));
+
+    smmu_vcmdq_write64(smmu_cmdqv_base, SMMU_VCMDQ_CMDQ_BASE,
+                       page_to_phys(parent_gpu->smmu_war.smmu_cmdq) | SMMU_VCMDQ_CMDQ_BASE_LOG2SIZE);
+    smmu_vcmdq_write32(smmu_cmdqv_base, SMMU_VCMDQ_CONS, 0);
+    smmu_vcmdq_write32(smmu_cmdqv_base, SMMU_VCMDQ_PROD, 0);
+    smmu_vcmdq_write32(smmu_cmdqv_base, SMMU_VCMDQ_CONFIG, SMMU_VCMDQ_CONFIG_ENABLE);
+    UVM_SPIN_WHILE(!(smmu_vcmdq_read32(smmu_cmdqv_base, SMMU_VCMDQ_STATUS) & SMMU_VCMDQ_STATUS_ENABLED), &spin);
+
+    uvm_mutex_init(&parent_gpu->smmu_war.smmu_lock, UVM_LOCK_ORDER_LEAF);
+    parent_gpu->smmu_war.smmu_prod = 0;
+    parent_gpu->smmu_war.smmu_cons = 0;
+
+    return NV_OK;
+
+out:
+    iounmap(parent_gpu->smmu_war.smmu_cmdqv_base);
+    parent_gpu->smmu_war.smmu_cmdqv_base = NULL;
+
+    return status;
+}
+
+static void uvm_ats_smmu_war_deinit(uvm_parent_gpu_t *parent_gpu)
+{
+    void __iomem *smmu_cmdqv_base = parent_gpu->smmu_war.smmu_cmdqv_base;
+    NvU32 cmdq_alloc_map;
+
+    if (parent_gpu->smmu_war.smmu_cmdqv_base) {
+        smmu_vcmdq_write32(smmu_cmdqv_base, SMMU_VCMDQ_CONFIG, 0);
+        cmdq_alloc_map = ioread32(smmu_cmdqv_base + SMMU_CMDQV_CMDQ_ALLOC_MAP(VCMDQ));
+        iowrite32(cmdq_alloc_map & SMMU_CMDQV_CMDQ_ALLOC_MAP_ALLOC, smmu_cmdqv_base + SMMU_CMDQV_CMDQ_ALLOC_MAP(VCMDQ));
+        smmu_vintf_write32(smmu_cmdqv_base, SMMU_VINTF_CONFIG, 0);
+    }
+
+    if (parent_gpu->smmu_war.smmu_cmdq)
+        __free_page(parent_gpu->smmu_war.smmu_cmdq);
+
+    if (parent_gpu->smmu_war.smmu_cmdqv_base)
+        iounmap(parent_gpu->smmu_war.smmu_cmdqv_base);
+}
+
+// The SMMU on ARM64 can run under different translation regimes depending on
+// what features the OS and CPU variant support. The CPU for GH180 supports
+// virtualisation extensions and starts the kernel at EL2 meaning SMMU operates
+// under the NS-EL2-E2H translation regime. Therefore we need to use the
+// TLBI_EL2_* commands which invalidate TLB entries created under this
+// translation regime.
+#define CMDQ_OP_TLBI_EL2_ASID 0x21;
+#define CMDQ_OP_TLBI_EL2_VA 0x22;
+#define CMDQ_OP_CMD_SYNC 0x46
+
+// Use the same maximum as used for MAX_TLBI_OPS in the upstream
+// kernel.
+#define UVM_MAX_TLBI_OPS (1UL << (PAGE_SHIFT - 3))
+
+#if UVM_ATS_SMMU_WAR_REQUIRED()
+void uvm_ats_smmu_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space, NvU64 addr, size_t size)
+{
+    struct mm_struct *mm = gpu_va_space->va_space->va_space_mm.mm;
+    uvm_parent_gpu_t *parent_gpu = gpu_va_space->gpu->parent;
+    struct {
+        NvU64 low;
+        NvU64 high;
+    } *vcmdq;
+    unsigned long vcmdq_prod;
+    NvU64 end;
+    uvm_spin_loop_t spin;
+    NvU16 asid;
+
+    if (!parent_gpu->smmu_war.smmu_cmdqv_base)
+        return;
+
+    asid = arm64_mm_context_get(mm);
+    vcmdq = kmap(parent_gpu->smmu_war.smmu_cmdq);
+    uvm_mutex_lock(&parent_gpu->smmu_war.smmu_lock);
+    vcmdq_prod = parent_gpu->smmu_war.smmu_prod;
+
+    // Our queue management is very simple. The mutex prevents multiple
+    // producers writing to the queue and all our commands require waiting for
+    // the queue to drain so we know it's empty. If we can't fit enough commands
+    // in the queue we just invalidate the whole ASID.
+    //
+    // The command queue is a cirular buffer with the MSB representing a wrap
+    // bit that must toggle on each wrap. See the SMMU architecture
+    // specification for more details.
+    //
+    // SMMU_VCMDQ_CMDQ_ENTRIES - 1 because we need to leave space for the
+    // CMD_SYNC.
+    if ((size >> PAGE_SHIFT) > min(UVM_MAX_TLBI_OPS, SMMU_VCMDQ_CMDQ_ENTRIES - 1)) {
+        vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].low = CMDQ_OP_TLBI_EL2_ASID;
+        vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].low |= (NvU64) asid << 48;
+        vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].high = 0;
+        vcmdq_prod++;
+    }
+    else {
+        for (end = addr + size; addr < end; addr += PAGE_SIZE) {
+            vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].low = CMDQ_OP_TLBI_EL2_VA;
+            vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].low |= (NvU64) asid << 48;
+            vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].high = addr & ~((1UL << 12) - 1);
+            vcmdq_prod++;
+        }
+    }
+
+    vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].low = CMDQ_OP_CMD_SYNC;
+    vcmdq[vcmdq_prod % SMMU_VCMDQ_CMDQ_ENTRIES].high = 0x0;
+    vcmdq_prod++;
+
+    // MSB is the wrap bit
+    vcmdq_prod &= (1UL << (SMMU_VCMDQ_CMDQ_BASE_LOG2SIZE + 1)) - 1;
+    parent_gpu->smmu_war.smmu_prod = vcmdq_prod;
+    smmu_vcmdq_write32(parent_gpu->smmu_war.smmu_cmdqv_base, SMMU_VCMDQ_PROD, parent_gpu->smmu_war.smmu_prod);
+
+    UVM_SPIN_WHILE(
+        (smmu_vcmdq_read32(parent_gpu->smmu_war.smmu_cmdqv_base, SMMU_VCMDQ_CONS) & GENMASK(19, 0)) != vcmdq_prod,
+        &spin);
+
+    uvm_mutex_unlock(&parent_gpu->smmu_war.smmu_lock);
+    kunmap(parent_gpu->smmu_war.smmu_cmdq);
+    arm64_mm_context_put(mm);
+}
+#endif
+
 NV_STATUS uvm_ats_sva_add_gpu(uvm_parent_gpu_t *parent_gpu)
 {
+#if NV_IS_EXPORT_SYMBOL_GPL_iommu_dev_enable_feature
     int ret;
 
     ret = iommu_dev_enable_feature(&parent_gpu->pci_dev->dev, IOMMU_DEV_FEAT_SVA);
-
-    return errno_to_nv_status(ret);
+    if (ret)
+        return errno_to_nv_status(ret);
+#endif
+    if (UVM_ATS_SMMU_WAR_REQUIRED())
+        return uvm_ats_smmu_war_init(parent_gpu);
+    else
+        return NV_OK;
 }
 
 void uvm_ats_sva_remove_gpu(uvm_parent_gpu_t *parent_gpu)
 {
+    if (UVM_ATS_SMMU_WAR_REQUIRED())
+        uvm_ats_smmu_war_deinit(parent_gpu);
+
+#if NV_IS_EXPORT_SYMBOL_GPL_iommu_dev_disable_feature
     iommu_dev_disable_feature(&parent_gpu->pci_dev->dev, IOMMU_DEV_FEAT_SVA);
+#endif
 }
 
 NV_STATUS uvm_ats_sva_bind_gpu(uvm_gpu_va_space_t *gpu_va_space)

kernel-open/nvidia-uvm/uvm_ats_sva.h

@@ -32,23 +32,38 @@
 // For ATS support on aarch64, arm_smmu_sva_bind() is needed for
 // iommu_sva_bind_device() calls. Unfortunately, arm_smmu_sva_bind() is not
 // conftest-able. We instead look for the presence of ioasid_get() or
-// mm_pasid_set(). ioasid_get() was added in the same patch series as
-// arm_smmu_sva_bind() and removed in v6.0. mm_pasid_set() was added in the
+// mm_pasid_drop(). ioasid_get() was added in the same patch series as
+// arm_smmu_sva_bind() and removed in v6.0. mm_pasid_drop() was added in the
 // same patch as the removal of ioasid_get(). We assume the presence of
-// arm_smmu_sva_bind() if ioasid_get(v5.11 - v5.17) or mm_pasid_set(v5.18+) is
+// arm_smmu_sva_bind() if ioasid_get(v5.11 - v5.17) or mm_pasid_drop(v5.18+) is
 // present.
 //
 // arm_smmu_sva_bind() was added with commit
 // 32784a9562fb0518b12e9797ee2aec52214adf6f and ioasid_get() was added with
 // commit cb4789b0d19ff231ce9f73376a023341300aed96 (11/23/2020). Commit
 // 701fac40384f07197b106136012804c3cae0b3de (02/15/2022) removed ioasid_get()
-// and added mm_pasid_set().
-    #if UVM_CAN_USE_MMU_NOTIFIERS() && (defined(NV_IOASID_GET_PRESENT) || defined(NV_MM_PASID_SET_PRESENT))
-        #define UVM_ATS_SVA_SUPPORTED() 1
+// and added mm_pasid_drop().
+    #if UVM_CAN_USE_MMU_NOTIFIERS() && (defined(NV_IOASID_GET_PRESENT) || defined(NV_MM_PASID_DROP_PRESENT))
+        #if defined(CONFIG_IOMMU_SVA)
+            #define UVM_ATS_SVA_SUPPORTED() 1
+        #else
+            #define UVM_ATS_SVA_SUPPORTED() 0
+        #endif
     #else
         #define UVM_ATS_SVA_SUPPORTED() 0
     #endif
 
+// If NV_ARCH_INVALIDATE_SECONDARY_TLBS is defined it means the upstream fix is
+// in place so no need for the WAR from Bug 4130089: [GH180][r535] WAR for
+// kernel not issuing SMMU TLB invalidates on read-only
+#if defined(NV_ARCH_INVALIDATE_SECONDARY_TLBS)
+    #define UVM_ATS_SMMU_WAR_REQUIRED() 0
+#elif NVCPU_IS_AARCH64
+    #define UVM_ATS_SMMU_WAR_REQUIRED() 1
+#else
+    #define UVM_ATS_SMMU_WAR_REQUIRED() 0
+#endif
+
 typedef struct
 {
     int placeholder;
@@ -77,6 +92,17 @@ typedef struct
 
     // LOCKING: None
     void uvm_ats_sva_unregister_gpu_va_space(uvm_gpu_va_space_t *gpu_va_space);
+
+    // Fix for Bug 4130089: [GH180][r535] WAR for kernel not issuing SMMU
+    // TLB invalidates on read-only to read-write upgrades
+    #if UVM_ATS_SMMU_WAR_REQUIRED()
+        void uvm_ats_smmu_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space, NvU64 addr, size_t size);
+    #else
+        static void uvm_ats_smmu_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space, NvU64 addr, size_t size)
+        {
+
+        }
+    #endif
 #else
     static NV_STATUS uvm_ats_sva_add_gpu(uvm_parent_gpu_t *parent_gpu)
     {
@@ -107,6 +133,11 @@ typedef struct
     {
 
     }
+
+    static void uvm_ats_smmu_invalidate_tlbs(uvm_gpu_va_space_t *gpu_va_space, NvU64 addr, size_t size)
+    {
+
+    }
 #endif // UVM_ATS_SVA_SUPPORTED
 
 #endif // __UVM_ATS_SVA_H__

kernel-open/nvidia-uvm/uvm_ce_test.c

@@ -191,7 +191,7 @@ static NV_STATUS test_membar(uvm_gpu_t *gpu)
 
     for (i = 0; i < REDUCTIONS; ++i) {
         uvm_push_set_flag(&push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE);
-        gpu->parent->ce_hal->semaphore_reduction_inc(&push, host_mem_gpu_va, REDUCTIONS + 1);
+        gpu->parent->ce_hal->semaphore_reduction_inc(&push, host_mem_gpu_va, REDUCTIONS);
     }
 
     // Without a sys membar the channel tracking semaphore can and does complete
@@ -577,7 +577,7 @@ static NV_STATUS test_semaphore_reduction_inc(uvm_gpu_t *gpu)
 
     for (i = 0; i < REDUCTIONS; i++) {
         uvm_push_set_flag(&push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE);
-        gpu->parent->ce_hal->semaphore_reduction_inc(&push, gpu_va, i+1);
+        gpu->parent->ce_hal->semaphore_reduction_inc(&push, gpu_va, REDUCTIONS);
     }
 
     status = uvm_push_end_and_wait(&push);
@@ -760,7 +760,7 @@ static NV_STATUS alloc_vidmem_protected(uvm_gpu_t *gpu, uvm_mem_t **mem, size_t
 
     *mem = NULL;
 
-    TEST_NV_CHECK_RET(uvm_mem_alloc_vidmem_protected(size, gpu, mem));
+    TEST_NV_CHECK_RET(uvm_mem_alloc_vidmem(size, gpu, mem));
     TEST_NV_CHECK_GOTO(uvm_mem_map_gpu_kernel(*mem, gpu), err);
     TEST_NV_CHECK_GOTO(zero_vidmem(*mem), err);
 

kernel-open/nvidia-uvm/uvm_channel.c

@@ -272,19 +272,26 @@ static bool try_claim_channel(uvm_channel_t *channel, NvU32 num_gpfifo_entries)
 
 static void unlock_channel_for_push(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_secure(channel)) {
-        NvU32 index = uvm_channel_index_in_pool(channel);
+    NvU32 index;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
-        uvm_channel_pool_assert_locked(channel->pool);
-        UVM_ASSERT(test_bit(index, channel->pool->push_locks));
-        __clear_bit(index, channel->pool->push_locks);
-        uvm_up_out_of_order(&channel->pool->push_sem);
-    }
+    if (!uvm_conf_computing_mode_enabled(gpu))
+        return;
+
+    index = uvm_channel_index_in_pool(channel);
+
+    uvm_channel_pool_assert_locked(channel->pool);
+    UVM_ASSERT(test_bit(index, channel->pool->push_locks));
+
+    __clear_bit(index, channel->pool->push_locks);
+    uvm_up_out_of_order(&channel->pool->push_sem);
 }
 
 static bool is_channel_locked_for_push(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_secure(channel))
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
+    if (uvm_conf_computing_mode_enabled(gpu))
         return test_bit(uvm_channel_index_in_pool(channel), channel->pool->push_locks);
 
     // For CE and proxy channels, we always return that the channel is locked,
@@ -295,25 +302,25 @@ static bool is_channel_locked_for_push(uvm_channel_t *channel)
 
 static void lock_channel_for_push(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_secure(channel)) {
-        NvU32 index = uvm_channel_index_in_pool(channel);
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+    NvU32 index = uvm_channel_index_in_pool(channel);
 
-        uvm_channel_pool_assert_locked(channel->pool);
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    uvm_channel_pool_assert_locked(channel->pool);
+    UVM_ASSERT(!test_bit(index, channel->pool->push_locks));
 
-        UVM_ASSERT(!test_bit(index, channel->pool->push_locks));
-        __set_bit(index, channel->pool->push_locks);
-    }
+    __set_bit(index, channel->pool->push_locks);
 }
 
 static bool test_claim_and_lock_channel(uvm_channel_t *channel, NvU32 num_gpfifo_entries)
 {
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     NvU32 index = uvm_channel_index_in_pool(channel);
 
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
     uvm_channel_pool_assert_locked(channel->pool);
 
-    if (uvm_channel_is_secure(channel) &&
-        !test_bit(index, channel->pool->push_locks) &&
-        try_claim_channel_locked(channel, num_gpfifo_entries)) {
+    if (!test_bit(index, channel->pool->push_locks) && try_claim_channel_locked(channel, num_gpfifo_entries)) {
         lock_channel_for_push(channel);
         return true;
     }
@@ -321,57 +328,15 @@ static bool test_claim_and_lock_channel(uvm_channel_t *channel, NvU32 num_gpfifo
     return false;
 }
 
-// Reserve a channel in the specified CE pool
-static NV_STATUS channel_reserve_in_ce_pool(uvm_channel_pool_t *pool, uvm_channel_t **channel_out)
-{
-    uvm_channel_t *channel;
-    uvm_spin_loop_t spin;
-
-    UVM_ASSERT(pool);
-    UVM_ASSERT(uvm_channel_pool_is_ce(pool));
-
-    uvm_for_each_channel_in_pool(channel, pool) {
-        // TODO: Bug 1764953: Prefer idle/less busy channels
-        if (try_claim_channel(channel, 1)) {
-            *channel_out = channel;
-            return NV_OK;
-        }
-    }
-
-    uvm_spin_loop_init(&spin);
-    while (1) {
-        uvm_for_each_channel_in_pool(channel, pool) {
-            NV_STATUS status;
-
-            uvm_channel_update_progress(channel);
-
-            if (try_claim_channel(channel, 1)) {
-                *channel_out = channel;
-
-                return NV_OK;
-            }
-
-            status = uvm_channel_check_errors(channel);
-            if (status != NV_OK)
-                return status;
-
-            UVM_SPIN_LOOP(&spin);
-        }
-    }
-
-    UVM_ASSERT_MSG(0, "Cannot get here?!\n");
-    return NV_ERR_GENERIC;
-}
-
-// Reserve a channel in the specified secure pool
-static NV_STATUS channel_reserve_in_secure_pool(uvm_channel_pool_t *pool, uvm_channel_t **channel_out)
+// Reserve a channel in the specified pool. The channel is locked until the push
+// ends
+static NV_STATUS channel_reserve_and_lock_in_pool(uvm_channel_pool_t *pool, uvm_channel_t **channel_out)
 {
     uvm_channel_t *channel;
     uvm_spin_loop_t spin;
     NvU32 index;
 
     UVM_ASSERT(pool);
-    UVM_ASSERT(pool->secure);
     UVM_ASSERT(uvm_conf_computing_mode_enabled(pool->manager->gpu));
 
     // This semaphore is uvm_up() in unlock_channel_for_push() as part of the
@@ -426,6 +391,51 @@ done:
     return NV_OK;
 }
 
+// Reserve a channel in the specified pool
+static NV_STATUS channel_reserve_in_pool(uvm_channel_pool_t *pool, uvm_channel_t **channel_out)
+{
+    uvm_channel_t *channel;
+    uvm_spin_loop_t spin;
+
+    UVM_ASSERT(pool);
+
+    if (uvm_conf_computing_mode_enabled(pool->manager->gpu))
+        return channel_reserve_and_lock_in_pool(pool, channel_out);
+
+    uvm_for_each_channel_in_pool(channel, pool) {
+        // TODO: Bug 1764953: Prefer idle/less busy channels
+        if (try_claim_channel(channel, 1)) {
+            *channel_out = channel;
+            return NV_OK;
+        }
+    }
+
+    uvm_spin_loop_init(&spin);
+    while (1) {
+        uvm_for_each_channel_in_pool(channel, pool) {
+            NV_STATUS status;
+
+            uvm_channel_update_progress(channel);
+
+            if (try_claim_channel(channel, 1)) {
+                *channel_out = channel;
+
+                return NV_OK;
+            }
+
+            status = uvm_channel_check_errors(channel);
+            if (status != NV_OK)
+                return status;
+
+            UVM_SPIN_LOOP(&spin);
+        }
+    }
+
+    UVM_ASSERT_MSG(0, "Cannot get here?!\n");
+
+    return NV_ERR_GENERIC;
+}
+
 NV_STATUS uvm_channel_reserve_type(uvm_channel_manager_t *manager, uvm_channel_type_t type, uvm_channel_t **channel_out)
 {
     uvm_channel_pool_t *pool = manager->pool_to_use.default_for_type[type];
@@ -433,10 +443,7 @@ NV_STATUS uvm_channel_reserve_type(uvm_channel_manager_t *manager, uvm_channel_t
     UVM_ASSERT(pool != NULL);
     UVM_ASSERT(type < UVM_CHANNEL_TYPE_COUNT);
 
-    if (pool->secure)
-        return channel_reserve_in_secure_pool(pool, channel_out);
-
-    return channel_reserve_in_ce_pool(pool, channel_out);
+    return channel_reserve_in_pool(pool, channel_out);
 }
 
 NV_STATUS uvm_channel_reserve_gpu_to_gpu(uvm_channel_manager_t *manager,
@@ -452,10 +459,7 @@ NV_STATUS uvm_channel_reserve_gpu_to_gpu(uvm_channel_manager_t *manager,
 
     UVM_ASSERT(pool->pool_type == UVM_CHANNEL_POOL_TYPE_CE);
 
-    if (pool->secure)
-        return channel_reserve_in_secure_pool(pool, channel_out);
-
-    return channel_reserve_in_ce_pool(pool, channel_out);
+    return channel_reserve_in_pool(pool, channel_out);
 }
 
 NV_STATUS uvm_channel_manager_wait(uvm_channel_manager_t *manager)
@@ -491,7 +495,7 @@ static NvU32 channel_get_available_push_info_index(uvm_channel_t *channel)
     return push_info - channel->push_infos;
 }
 
-static void channel_semaphore_gpu_encrypt_payload(uvm_push_t *push, uvm_channel_t *channel, NvU64 semaphore_va)
+static void channel_semaphore_gpu_encrypt_payload(uvm_push_t *push, NvU64 semaphore_va)
 {
     NvU32 iv_index;
     uvm_gpu_address_t notifier_gpu_va;
@@ -499,12 +503,14 @@ static void channel_semaphore_gpu_encrypt_payload(uvm_push_t *push, uvm_channel_
     uvm_gpu_address_t semaphore_gpu_va;
     uvm_gpu_address_t encrypted_payload_gpu_va;
     uvm_gpu_t *gpu = push->gpu;
+    uvm_channel_t *channel = push->channel;
     uvm_gpu_semaphore_t *semaphore = &channel->tracking_sem.semaphore;
     UvmCslIv *iv_cpu_addr = semaphore->conf_computing.ivs;
     NvU32 payload_size = sizeof(*semaphore->payload);
     NvU32 *last_pushed_notifier = &semaphore->conf_computing.last_pushed_notifier;
 
-    UVM_ASSERT(uvm_channel_is_secure_ce(channel));
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
 
     encrypted_payload_gpu_va = uvm_rm_mem_get_gpu_va(semaphore->conf_computing.encrypted_payload, gpu, false);
     notifier_gpu_va = uvm_rm_mem_get_gpu_va(semaphore->conf_computing.notifier, gpu, false);
@@ -538,19 +544,21 @@ NV_STATUS uvm_channel_begin_push(uvm_channel_t *channel, uvm_push_t *push)
 {
     NV_STATUS status;
     uvm_channel_manager_t *manager;
+    uvm_gpu_t *gpu;
 
     UVM_ASSERT(channel);
     UVM_ASSERT(push);
 
     manager = channel->pool->manager;
 
+    gpu = uvm_channel_get_gpu(channel);
+
     // Only SEC2 and WLC with set up fixed schedule can use direct push
     // submission. All other cases (including WLC pre-schedule) need to
     // reserve a launch channel that will be used to submit this push
     // indirectly.
-    if (uvm_conf_computing_mode_enabled(uvm_channel_get_gpu(channel)) &&
-        !(uvm_channel_is_wlc(channel) && uvm_channel_manager_is_wlc_ready(manager)) &&
-        !uvm_channel_is_sec2(channel)) {
+    if (uvm_conf_computing_mode_enabled(gpu) && uvm_channel_is_ce(channel) &&
+        !(uvm_channel_is_wlc(channel) && uvm_channel_manager_is_wlc_ready(manager))) {
         uvm_channel_type_t indirect_channel_type = uvm_channel_manager_is_wlc_ready(manager) ?
                                                    UVM_CHANNEL_TYPE_WLC :
                                                    UVM_CHANNEL_TYPE_SEC2;
@@ -559,9 +567,9 @@ NV_STATUS uvm_channel_begin_push(uvm_channel_t *channel, uvm_push_t *push)
             return status;
     }
 
-    // For secure channels, channel's lock should have been acquired in
-    // uvm_channel_reserve() or channel_reserve_in_secure_pool() before
-    // reaching here.
+    // When the Confidential Computing feature is enabled, the channel's lock
+    // should have already been acquired in uvm_channel_reserve() or
+    // channel_reserve_and_lock_in_pool().
     UVM_ASSERT(is_channel_locked_for_push(channel));
 
     push->channel = channel;
@@ -586,9 +594,8 @@ static void internal_channel_submit_work(uvm_push_t *push, NvU32 push_size, NvU3
     NvU64 *gpfifo_entry;
     NvU64 pushbuffer_va;
     uvm_channel_t *channel = push->channel;
-    uvm_channel_manager_t *channel_manager = channel->pool->manager;
-    uvm_pushbuffer_t *pushbuffer = channel_manager->pushbuffer;
-    uvm_gpu_t *gpu = channel_manager->gpu;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
     BUILD_BUG_ON(sizeof(*gpfifo_entry) != NVB06F_GP_ENTRY__SIZE);
     UVM_ASSERT(!uvm_channel_is_proxy(channel));
@@ -644,12 +651,11 @@ static void proxy_channel_submit_work(uvm_push_t *push, NvU32 push_size)
 static void do_semaphore_release(uvm_push_t *push, NvU64 semaphore_va, NvU32 new_payload)
 {
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
+
     if (uvm_channel_is_ce(push->channel))
         gpu->parent->ce_hal->semaphore_release(push, semaphore_va, new_payload);
-    else if (uvm_channel_is_sec2(push->channel))
-        gpu->parent->sec2_hal->semaphore_release(push, semaphore_va, new_payload);
     else
-        UVM_ASSERT_MSG(0, "Semaphore release on an unsupported channel.\n");
+        gpu->parent->sec2_hal->semaphore_release(push, semaphore_va, new_payload);
 }
 
 static void uvm_channel_tracking_semaphore_release(uvm_push_t *push, NvU64 semaphore_va, NvU32 new_payload)
@@ -668,8 +674,8 @@ static void uvm_channel_tracking_semaphore_release(uvm_push_t *push, NvU64 semap
     // needs to be scheduled to get an encrypted shadow copy in unprotected
     // sysmem. This allows UVM to later decrypt it and observe the new
     // semaphore value.
-    if (uvm_channel_is_secure_ce(push->channel))
-        channel_semaphore_gpu_encrypt_payload(push, push->channel, semaphore_va);
+    if (uvm_conf_computing_mode_enabled(push->gpu) && uvm_channel_is_ce(push->channel))
+        channel_semaphore_gpu_encrypt_payload(push, semaphore_va);
 }
 
 static uvm_channel_t *get_paired_channel(uvm_channel_t *channel)
@@ -716,7 +722,17 @@ static void internal_channel_submit_work_wlc(uvm_push_t *push)
 
     // Wait for the WLC/LCIC to be primed. This means that PUT == GET + 2
     // and a WLC doorbell ring is enough to start work.
-    UVM_SPIN_WHILE(!uvm_gpu_tracking_semaphore_is_completed(&lcic_channel->tracking_sem), &spin);
+    UVM_SPIN_WHILE(!uvm_gpu_tracking_semaphore_is_completed(&lcic_channel->tracking_sem), &spin) {
+        NV_STATUS status = uvm_channel_check_errors(lcic_channel);
+        if (status != NV_OK) {
+            UVM_ASSERT(uvm_global_get_status() != NV_OK);
+
+            // If there's a global fatal error we can't communicate with the GPU
+            // and the below launch sequence doesn't work.
+            UVM_ERR_PRINT_NV_STATUS("Failed to wait for LCIC channel (%s) completion.", status, lcic_channel->name);
+            return;
+        }
+    }
 
     // Executing WLC adds an extra job to LCIC
     ++lcic_channel->tracking_sem.queued_value;
@@ -746,15 +762,12 @@ static void internal_channel_submit_work_wlc(uvm_push_t *push)
     wmb();
 
     // Ring the WLC doorbell to start processing the above push
-    UVM_GPU_WRITE_ONCE(*wlc_channel->channel_info.workSubmissionOffset,
-                       wlc_channel->channel_info.workSubmissionToken);
+    UVM_GPU_WRITE_ONCE(*wlc_channel->channel_info.workSubmissionOffset, wlc_channel->channel_info.workSubmissionToken);
 }
 
-static void internal_channel_submit_work_indirect_wlc(uvm_push_t *push,
-                                                      NvU32 old_cpu_put,
-                                                      NvU32 new_gpu_put)
+static void internal_channel_submit_work_indirect_wlc(uvm_push_t *push, NvU32 old_cpu_put, NvU32 new_gpu_put)
 {
-    uvm_pushbuffer_t *pushbuffer = push->channel->pool->manager->pushbuffer;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(push->channel);
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
 
     uvm_push_t indirect_push;
@@ -767,7 +780,7 @@ static void internal_channel_submit_work_indirect_wlc(uvm_push_t *push,
     uvm_gpu_address_t push_enc_auth_tag_gpu;
     NvU64 gpfifo_gpu_va = push->channel->channel_info.gpFifoGpuVa + old_cpu_put * sizeof(gpfifo_entry);
 
-    UVM_ASSERT(!uvm_channel_is_sec2(push->channel));
+    UVM_ASSERT(uvm_channel_is_ce(push->channel));
     UVM_ASSERT(uvm_channel_is_wlc(push->launch_channel));
 
     // WLC submissions are done under channel lock, so there should be no
@@ -848,8 +861,6 @@ static void update_gpput_via_sec2(uvm_push_t *sec2_push, uvm_channel_t *channel,
                                                            UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT,
                                                            &gpput_auth_tag_gpu);
 
-
-
     // Update GPPUT. The update needs 4B write to specific offset,
     // however we can only do 16B aligned decrypt writes.
     // A poison value is written to all other locations, this is ignored in
@@ -922,7 +933,7 @@ static void set_gpfifo_via_sec2(uvm_push_t *sec2_push, uvm_channel_t *channel, N
             gpfifo_scratchpad[0] = previous_gpfifo->control_value;
         }
         else {
-            uvm_pushbuffer_t *pushbuffer = channel->pool->manager->pushbuffer;
+            uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
             NvU64 prev_pb_va = uvm_pushbuffer_get_gpu_va_base(pushbuffer) + previous_gpfifo->pushbuffer_offset;
 
             // Reconstruct the previous gpfifo entry. UVM_GPFIFO_SYNC_WAIT is
@@ -951,11 +962,9 @@ static void set_gpfifo_via_sec2(uvm_push_t *sec2_push, uvm_channel_t *channel, N
                                    gpfifo_auth_tag_gpu.address);
 }
 
-static NV_STATUS internal_channel_submit_work_indirect_sec2(uvm_push_t *push,
-                                                            NvU32 old_cpu_put,
-                                                            NvU32 new_gpu_put)
+static NV_STATUS internal_channel_submit_work_indirect_sec2(uvm_push_t *push, NvU32 old_cpu_put, NvU32 new_gpu_put)
 {
-    uvm_pushbuffer_t *pushbuffer = push->channel->pool->manager->pushbuffer;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(push->channel);
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
 
     uvm_push_t indirect_push;
@@ -968,7 +977,7 @@ static NV_STATUS internal_channel_submit_work_indirect_sec2(uvm_push_t *push,
     uvm_gpu_address_t push_auth_tag_gpu;
     uvm_spin_loop_t spin;
 
-    UVM_ASSERT(!uvm_channel_is_sec2(push->channel));
+    UVM_ASSERT(uvm_channel_is_ce(push->channel));
     UVM_ASSERT(uvm_channel_is_sec2(push->launch_channel));
 
     // If the old_cpu_put is not equal to the last gpu put, other pushes are
@@ -1051,7 +1060,7 @@ static void encrypt_push(uvm_push_t *push)
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
     NvU32 push_size = uvm_push_get_size(push);
     uvm_push_info_t *push_info = uvm_push_info_from_push(push);
-    uvm_pushbuffer_t *pushbuffer = channel->pool->manager->pushbuffer;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
     unsigned auth_tag_offset = UVM_CONF_COMPUTING_AUTH_TAG_SIZE * push->push_info_index;
 
     if (!uvm_conf_computing_mode_enabled(gpu))
@@ -1098,6 +1107,7 @@ void uvm_channel_end_push(uvm_push_t *push)
     NvU32 push_size;
     NvU32 cpu_put;
     NvU32 new_cpu_put;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     bool needs_sec2_work_submit = false;
 
     channel_pool_lock(channel->pool);
@@ -1112,7 +1122,7 @@ void uvm_channel_end_push(uvm_push_t *push)
 
     if (uvm_channel_is_wlc(channel) && uvm_channel_manager_is_wlc_ready(channel_manager)) {
         uvm_channel_t *paired_lcic = wlc_get_paired_lcic(channel);
-        uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
         gpu->parent->ce_hal->semaphore_reduction_inc(push,
                                                      paired_lcic->channel_info.gpPutGpuVa,
                                                      paired_lcic->num_gpfifo_entries - 1);
@@ -1126,7 +1136,7 @@ void uvm_channel_end_push(uvm_push_t *push)
             // pushes. However, direct pushes to WLC can be smaller than this
             // size. This is used e.g. by indirect submission of control
             // gpfifo entries.
-            channel_manager->gpu->parent->host_hal->noop(push, UVM_MAX_WLC_PUSH_SIZE - uvm_push_get_size(push));
+            gpu->parent->host_hal->noop(push, UVM_MAX_WLC_PUSH_SIZE - uvm_push_get_size(push));
         }
     }
 
@@ -1144,7 +1154,7 @@ void uvm_channel_end_push(uvm_push_t *push)
     // Indirect submission via SEC2/WLC needs pushes to be aligned for
     // encryption/decryption. The pushbuffer_size of this push
     // influences starting address of the next push.
-    if (uvm_conf_computing_mode_enabled(uvm_channel_get_gpu(channel)))
+    if (uvm_conf_computing_mode_enabled(gpu))
         entry->pushbuffer_size = UVM_ALIGN_UP(push_size, UVM_CONF_COMPUTING_BUF_ALIGNMENT);
     entry->push_info = &channel->push_infos[push->push_info_index];
     entry->type = UVM_GPFIFO_ENTRY_TYPE_NORMAL;
@@ -1158,12 +1168,13 @@ void uvm_channel_end_push(uvm_push_t *push)
     else if (uvm_channel_is_wlc(channel) && uvm_channel_manager_is_wlc_ready(channel_manager)) {
         internal_channel_submit_work_wlc(push);
     }
-    else if (uvm_conf_computing_mode_enabled(channel_manager->gpu) && !uvm_channel_is_sec2(channel)) {
+    else if (uvm_conf_computing_mode_enabled(gpu) && uvm_channel_is_ce(channel)) {
         if (uvm_channel_manager_is_wlc_ready(channel_manager)) {
             internal_channel_submit_work_indirect_wlc(push, cpu_put, new_cpu_put);
         }
         else {
-            // submitting via SEC2 starts a push, postpone until this push is ended
+            // submitting via SEC2 starts a push, postpone until this push is
+            // ended
             needs_sec2_work_submit = true;
         }
     }
@@ -1202,12 +1213,13 @@ void uvm_channel_end_push(uvm_push_t *push)
 
 static void submit_ctrl_gpfifo(uvm_channel_t *channel, uvm_gpfifo_entry_t *entry, NvU32 new_cpu_put)
 {
-    uvm_gpu_t *gpu = channel->pool->manager->gpu;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     NvU32 cpu_put = channel->cpu_put;
     NvU64 *gpfifo_entry;
 
     UVM_ASSERT(entry == &channel->gpfifo_entries[cpu_put]);
-    if (uvm_conf_computing_mode_enabled(gpu) && !uvm_channel_is_sec2(channel))
+
+    if (uvm_conf_computing_mode_enabled(gpu) && uvm_channel_is_ce(channel))
         return;
 
     gpfifo_entry = (NvU64*)channel->channel_info.gpFifoEntries + cpu_put;
@@ -1234,7 +1246,7 @@ static NV_STATUS submit_ctrl_gpfifo_indirect(uvm_channel_t *channel,
                                                UVM_CHANNEL_TYPE_WLC :
                                                UVM_CHANNEL_TYPE_SEC2;
 
-    UVM_ASSERT(!uvm_channel_is_sec2(channel));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
 
     // If the old_cpu_put is not equal to the last gpu put,
     // Another push(es) is pending that needs to be submitted.
@@ -1290,6 +1302,7 @@ static void write_ctrl_gpfifo(uvm_channel_t *channel, NvU64 ctrl_fifo_entry_valu
     NvU32 cpu_put;
     NvU32 new_cpu_put;
     bool needs_indirect_submit = false;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
     channel_pool_lock(channel->pool);
 
@@ -1312,7 +1325,7 @@ static void write_ctrl_gpfifo(uvm_channel_t *channel, NvU64 ctrl_fifo_entry_valu
     --channel->current_gpfifo_count;
 
     submit_ctrl_gpfifo(channel, entry, new_cpu_put);
-    if (uvm_conf_computing_mode_enabled(channel->pool->manager->gpu) && !uvm_channel_is_sec2(channel))
+    if (uvm_conf_computing_mode_enabled(gpu) && uvm_channel_is_ce(channel))
         needs_indirect_submit = true;
 
     channel->cpu_put = new_cpu_put;
@@ -1385,16 +1398,15 @@ NV_STATUS uvm_channel_write_ctrl_gpfifo(uvm_channel_t *channel, NvU64 ctrl_fifo_
     return NV_OK;
 }
 
-static NV_STATUS uvm_channel_reserve_secure(uvm_channel_t *channel, NvU32 num_gpfifo_entries)
+static NV_STATUS channel_reserve_and_lock(uvm_channel_t *channel, NvU32 num_gpfifo_entries)
 {
     uvm_spin_loop_t spin;
-    NV_STATUS status = NV_OK;
     uvm_channel_pool_t *pool = channel->pool;
 
     // This semaphore is uvm_up() in unlock_channel_for_push() as part of the
     // uvm_channel_end_push() routine. Note that different than in
-    // channel_reserve_in_secure_pool, we cannot pick an unlocked channel from
-    // the secure pool, even when there is one available and *channel is locked.
+    // channel_reserve_and_lock_in_pool, we cannot pick an unlocked channel from
+    // the pool, even when there is one available and *channel is locked.
     // Not a concern given that uvm_channel_reserve() is not the common-case for
     // channel reservation, and only used for channel initialization, GPFIFO
     // control work submission, and testing.
@@ -1409,6 +1421,8 @@ static NV_STATUS uvm_channel_reserve_secure(uvm_channel_t *channel, NvU32 num_gp
 
     uvm_spin_loop_init(&spin);
     while (1) {
+        NV_STATUS status;
+
         uvm_channel_update_progress(channel);
 
         channel_pool_lock(pool);
@@ -1436,9 +1450,10 @@ NV_STATUS uvm_channel_reserve(uvm_channel_t *channel, NvU32 num_gpfifo_entries)
 {
     NV_STATUS status = NV_OK;
     uvm_spin_loop_t spin;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
-    if (uvm_channel_is_secure(channel))
-        return uvm_channel_reserve_secure(channel, num_gpfifo_entries);
+    if (uvm_conf_computing_mode_enabled(gpu))
+        return channel_reserve_and_lock(channel, num_gpfifo_entries);
 
     if (try_claim_channel(channel, num_gpfifo_entries))
         return NV_OK;
@@ -1578,8 +1593,10 @@ NvU64 uvm_channel_update_completed_value(uvm_channel_t *channel)
 static NV_STATUS csl_init(uvm_channel_t *channel)
 {
     NV_STATUS status;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
 
-    UVM_ASSERT(uvm_channel_is_secure(channel));
     uvm_mutex_init(&channel->csl.ctx_lock, UVM_LOCK_ORDER_LEAF);
 
     status = uvm_rm_locked_call(nvUvmInterfaceCslInitContext(&channel->csl.ctx, channel->handle));
@@ -1589,7 +1606,7 @@ static NV_STATUS csl_init(uvm_channel_t *channel)
     else {
         UVM_DBG_PRINT("nvUvmInterfaceCslInitContext() failed: %s, GPU %s\n",
                       nvstatusToString(status),
-                      uvm_gpu_name(channel->pool->manager->gpu));
+                      uvm_gpu_name(gpu));
     }
 
     return status;
@@ -1609,7 +1626,10 @@ static void csl_destroy(uvm_channel_t *channel)
 
 static void free_conf_computing_buffers(uvm_channel_t *channel)
 {
-    UVM_ASSERT(uvm_channel_is_secure_ce(channel));
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
 
     uvm_rm_mem_free(channel->conf_computing.static_pb_protected_vidmem);
     uvm_rm_mem_free(channel->conf_computing.static_pb_unprotected_sysmem);
@@ -1637,10 +1657,12 @@ static void free_conf_computing_buffers(uvm_channel_t *channel)
 static NV_STATUS alloc_conf_computing_buffers_semaphore(uvm_channel_t *channel)
 {
     uvm_gpu_semaphore_t *semaphore = &channel->tracking_sem.semaphore;
-    uvm_gpu_t *gpu = channel->pool->manager->gpu;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     NV_STATUS status;
 
-    UVM_ASSERT(uvm_channel_is_secure_ce(channel));
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
+
     status = uvm_rm_mem_alloc_and_map_cpu(gpu,
                                           UVM_RM_MEM_TYPE_SYS,
                                           sizeof(semaphore->conf_computing.last_pushed_notifier),
@@ -1679,7 +1701,7 @@ static NV_STATUS alloc_conf_computing_buffers_semaphore(uvm_channel_t *channel)
 
 static NV_STATUS alloc_conf_computing_buffers_wlc(uvm_channel_t *channel)
 {
-    uvm_gpu_t *gpu = channel->pool->manager->gpu;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     size_t aligned_wlc_push_size = UVM_ALIGN_UP(UVM_MAX_WLC_PUSH_SIZE, UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT);
     NV_STATUS status = uvm_rm_mem_alloc_and_map_cpu(gpu,
                                                     UVM_RM_MEM_TYPE_SYS,
@@ -1723,7 +1745,7 @@ static NV_STATUS alloc_conf_computing_buffers_wlc(uvm_channel_t *channel)
 
 static NV_STATUS alloc_conf_computing_buffers_lcic(uvm_channel_t *channel)
 {
-    uvm_gpu_t *gpu = channel->pool->manager->gpu;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
     const size_t notifier_size = sizeof(*channel->conf_computing.static_notifier_entry_unprotected_sysmem_cpu);
     NV_STATUS status = uvm_rm_mem_alloc_and_map_cpu(gpu,
                                                     UVM_RM_MEM_TYPE_SYS,
@@ -1758,8 +1780,10 @@ static NV_STATUS alloc_conf_computing_buffers_lcic(uvm_channel_t *channel)
 static NV_STATUS alloc_conf_computing_buffers(uvm_channel_t *channel)
 {
     NV_STATUS status;
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
-    UVM_ASSERT(uvm_channel_is_secure_ce(channel));
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
 
     status = alloc_conf_computing_buffers_semaphore(channel);
     if (status != NV_OK)
@@ -1772,7 +1796,6 @@ static NV_STATUS alloc_conf_computing_buffers(uvm_channel_t *channel)
         status = alloc_conf_computing_buffers_lcic(channel);
     }
     else {
-        uvm_gpu_t *gpu = channel->pool->manager->gpu;
         void *push_crypto_bundles = uvm_kvmalloc_zero(sizeof(*channel->conf_computing.push_crypto_bundles) *
                                                       channel->num_gpfifo_entries);
 
@@ -1793,6 +1816,8 @@ static NV_STATUS alloc_conf_computing_buffers(uvm_channel_t *channel)
 
 static void channel_destroy(uvm_channel_pool_t *pool, uvm_channel_t *channel)
 {
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
     UVM_ASSERT(pool->num_channels > 0);
 
     if (channel->tracking_sem.queued_value > 0) {
@@ -1816,9 +1841,10 @@ static void channel_destroy(uvm_channel_pool_t *pool, uvm_channel_t *channel)
 
     uvm_kvfree(channel->gpfifo_entries);
 
-    if (uvm_channel_is_secure(channel)) {
+    if (uvm_conf_computing_mode_enabled(gpu)) {
         csl_destroy(channel);
-        if (uvm_channel_is_secure_ce(channel))
+
+        if (uvm_channel_is_ce(channel))
             free_conf_computing_buffers(channel);
     }
 
@@ -1905,8 +1931,6 @@ static NV_STATUS internal_channel_create(uvm_channel_t *channel)
         channel_alloc_params.gpPutLoc = UVM_BUFFER_LOCATION_SYS;
     }
 
-    channel_alloc_params.secure = channel->pool->secure;
-
     status = uvm_rm_locked_call(nvUvmInterfaceChannelAllocate(channel_get_tsg(channel),
                                                               &channel_alloc_params,
                                                               &channel->handle,
@@ -1928,8 +1952,7 @@ static NV_STATUS internal_channel_create(uvm_channel_t *channel)
              channel_info->hwChannelId,
              uvm_channel_is_sec2(channel) ? "SEC2" :
              uvm_channel_is_wlc(channel) ? "WLC" :
-             uvm_channel_is_lcic(channel) ? "LCIC" :
-             uvm_channel_is_secure(channel) ? "CE (secure)" : "CE",
+             uvm_channel_is_lcic(channel) ? "LCIC" : "CE",
              channel->pool->engine_index);
 
     return NV_OK;
@@ -1981,7 +2004,7 @@ static NV_STATUS channel_create(uvm_channel_pool_t *pool, uvm_channel_t *channel
     channel->tools.pending_event_count = 0;
     INIT_LIST_HEAD(&channel->tools.channel_list_node);
 
-    if (uvm_conf_computing_mode_enabled(gpu) && !uvm_channel_is_sec2(channel))
+    if (uvm_conf_computing_mode_enabled(gpu) && uvm_channel_is_ce(channel))
         semaphore_pool = gpu->secure_semaphore_pool;
 
     status = uvm_gpu_tracking_semaphore_alloc(semaphore_pool, &channel->tracking_sem);
@@ -2007,7 +2030,7 @@ static NV_STATUS channel_create(uvm_channel_pool_t *pool, uvm_channel_t *channel
         goto error;
     }
 
-    if (uvm_channel_is_secure(channel)) {
+    if (uvm_conf_computing_mode_enabled(gpu)) {
         status = csl_init(channel);
         if (status != NV_OK)
             goto error;
@@ -2075,7 +2098,7 @@ static NV_STATUS channel_init(uvm_channel_t *channel)
 
     if (uvm_gpu_has_pushbuffer_segments(gpu)) {
         NvU64 gpfifo_entry;
-        uvm_pushbuffer_t *pushbuffer = channel->pool->manager->pushbuffer;
+        uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
         NvU64 pb_base = uvm_pushbuffer_get_gpu_va_base(pushbuffer);
 
         if (uvm_channel_is_sec2(channel))
@@ -2095,10 +2118,8 @@ static NV_STATUS channel_init(uvm_channel_t *channel)
 
     if (uvm_channel_is_ce(channel))
         gpu->parent->ce_hal->init(&push);
-    else if (uvm_channel_is_sec2(channel))
-        gpu->parent->sec2_hal->init(&push);
     else
-        UVM_ASSERT_MSG(0, "Unknown channel type!");
+        gpu->parent->sec2_hal->init(&push);
 
     gpu->parent->host_hal->init(&push);
 
@@ -2153,11 +2174,6 @@ static unsigned channel_pool_type_num_tsgs(uvm_channel_pool_type_t pool_type)
     return 1;
 }
 
-static bool pool_type_is_valid(uvm_channel_pool_type_t pool_type)
-{
-    return(is_power_of_2(pool_type) && (pool_type < UVM_CHANNEL_POOL_TYPE_MASK));
-}
-
 static UVM_GPU_CHANNEL_ENGINE_TYPE pool_type_to_engine_type(uvm_channel_pool_type_t pool_type)
 {
     if (pool_type ==  UVM_CHANNEL_POOL_TYPE_SEC2)
@@ -2229,7 +2245,7 @@ static NV_STATUS channel_pool_add(uvm_channel_manager_t *channel_manager,
     unsigned num_tsgs;
     uvm_channel_pool_t *pool;
 
-    UVM_ASSERT(pool_type_is_valid(pool_type));
+    UVM_ASSERT(uvm_pool_type_is_valid(pool_type));
 
     pool = channel_manager->channel_pools + channel_manager->num_channel_pools;
     channel_manager->num_channel_pools++;
@@ -2260,10 +2276,10 @@ static NV_STATUS channel_pool_add(uvm_channel_manager_t *channel_manager,
     num_channels = channel_pool_type_num_channels(pool_type);
     UVM_ASSERT(num_channels <= UVM_CHANNEL_MAX_NUM_CHANNELS_PER_POOL);
 
-    if (pool->secure) {
+    if (uvm_conf_computing_mode_enabled(channel_manager->gpu)) {
         // Use different order lock for SEC2 and WLC channels.
         // This allows reserving a SEC2 or WLC channel for indirect work
-        // submission while holding a reservation for a secure channel.
+        // submission while holding a reservation for a channel.
         uvm_lock_order_t order = uvm_channel_pool_is_sec2(pool) ? UVM_LOCK_ORDER_CSL_SEC2_PUSH :
                                  (uvm_channel_pool_is_wlc(pool) ? UVM_LOCK_ORDER_CSL_WLC_PUSH :
                                                                   UVM_LOCK_ORDER_CSL_PUSH);
@@ -2297,23 +2313,6 @@ static NV_STATUS channel_pool_add(uvm_channel_manager_t *channel_manager,
     return status;
 }
 
-static NV_STATUS channel_pool_add_secure(uvm_channel_manager_t *channel_manager,
-                                         uvm_channel_pool_type_t pool_type,
-                                         unsigned engine_index,
-                                         uvm_channel_pool_t **pool_out)
-{
-    uvm_channel_pool_t *pool = channel_manager->channel_pools + channel_manager->num_channel_pools;
-
-    pool->secure = true;
-    return channel_pool_add(channel_manager, pool_type, engine_index, pool_out);
-}
-
-bool uvm_channel_type_requires_secure_pool(uvm_gpu_t *gpu, uvm_channel_type_t channel_type)
-{
-    // For now, all channels are secure channels
-    return true;
-}
-
 static bool ce_usable_for_channel_type(uvm_channel_type_t type, const UvmGpuCopyEngineCaps *cap)
 {
     if (!cap->supported || cap->grce)
@@ -2461,13 +2460,6 @@ static NV_STATUS pick_ce_for_channel_type(uvm_channel_manager_t *manager,
         if (!ce_usable_for_channel_type(type, cap))
             continue;
 
-        if (uvm_conf_computing_mode_is_hcc(manager->gpu)) {
-            // All usable CEs are secure
-            UVM_ASSERT(cap->secure);
-
-            // Multi-PCE LCEs are disallowed
-            UVM_ASSERT(hweight32(cap->cePceMask) == 1);
-        }
         __set_bit(i, manager->ce_mask);
 
         if (best_ce == UVM_COPY_ENGINE_COUNT_MAX) {
@@ -2523,7 +2515,7 @@ out:
     return status;
 }
 
-// Return the non-secure pool corresponding to the given CE index
+// Return the pool corresponding to the given CE index
 //
 // This function cannot be used to access the proxy pool in SR-IOV heavy.
 static uvm_channel_pool_t *channel_manager_ce_pool(uvm_channel_manager_t *manager, NvU32 ce)
@@ -2703,7 +2695,15 @@ static void init_channel_manager_conf(uvm_channel_manager_t *manager)
     // access through the bus, because no cache coherence message is exchanged.
     if (uvm_gpu_is_coherent(gpu->parent)) {
         manager->conf.gpfifo_loc = UVM_BUFFER_LOCATION_SYS;
-        manager->conf.gpput_loc = UVM_BUFFER_LOCATION_SYS;
+
+        // On GPUs with limited ESCHED addressing range, e.g., Volta on P9, RM
+        // cannot guarantee that USERD/GPPUT physical address is accessible by
+        // ESCHED. We set GPPUT location to vidmem where physical addresses are
+        // all accessible by ESCHED. We use the max_host_va as a proxy for the
+        // PA limitation, since all architectures with 40b VA limits also have
+        // 40b PA limits.
+        manager->conf.gpput_loc = gpu->parent->max_host_va == (1ull << 40) ? UVM_BUFFER_LOCATION_VID :
+                                                                             UVM_BUFFER_LOCATION_SYS;
     }
     else {
         // By default we place GPFIFO and GPPUT on vidmem as it potentially has
@@ -2726,24 +2726,17 @@ static void init_channel_manager_conf(uvm_channel_manager_t *manager)
 static unsigned channel_manager_get_max_pools(uvm_channel_manager_t *manager)
 {
     unsigned num_channel_pools;
-    unsigned num_used_ce = bitmap_weight(manager->ce_mask, UVM_COPY_ENGINE_COUNT_MAX);
 
     // Create one CE channel pool per usable CE
-    num_channel_pools = num_used_ce;
+    num_channel_pools = bitmap_weight(manager->ce_mask, UVM_COPY_ENGINE_COUNT_MAX);
 
     // CE proxy channel pool.
     if (uvm_gpu_uses_proxy_channel_pool(manager->gpu))
         num_channel_pools++;
 
-    if (uvm_conf_computing_mode_enabled(manager->gpu)) {
-
-        // Create one CE secure channel pool per usable CE
-        if (uvm_conf_computing_mode_is_hcc(manager->gpu))
-            num_channel_pools += num_used_ce;
-
-        // SEC2 pool, WLC pool, LCIC pool
+    // SEC2 pool, WLC pool, LCIC pool
+    if (uvm_conf_computing_mode_enabled(manager->gpu))
         num_channel_pools += 3;
-    }
 
     return num_channel_pools;
 }
@@ -2775,38 +2768,6 @@ static NV_STATUS channel_manager_create_ce_pools(uvm_channel_manager_t *manager,
     return NV_OK;
 }
 
-static NV_STATUS channel_manager_create_ce_secure_pools(uvm_channel_manager_t *manager, unsigned *preferred_ce)
-{
-    unsigned ce;
-
-    if (!uvm_conf_computing_mode_is_hcc(manager->gpu))
-        return NV_OK;
-
-    for_each_set_bit(ce, manager->ce_mask, UVM_COPY_ENGINE_COUNT_MAX) {
-        NV_STATUS status;
-        unsigned type;
-        uvm_channel_pool_t *pool = NULL;
-
-        status = channel_pool_add_secure(manager, UVM_CHANNEL_POOL_TYPE_CE, ce, &pool);
-        if (status != NV_OK)
-            return status;
-
-        for (type = 0; type < UVM_CHANNEL_TYPE_CE_COUNT; type++) {
-            unsigned preferred = preferred_ce[type];
-
-            if (preferred != ce)
-                continue;
-
-            if (uvm_channel_type_requires_secure_pool(manager->gpu, type)) {
-                UVM_ASSERT(manager->pool_to_use.default_for_type[type] == NULL);
-                manager->pool_to_use.default_for_type[type] = pool;
-            }
-        }
-    }
-
-    return NV_OK;
-}
-
 static NV_STATUS setup_wlc_schedule(uvm_channel_t *wlc)
 {
     uvm_gpu_t *gpu = uvm_channel_get_gpu(wlc);
@@ -3134,6 +3095,64 @@ static NV_STATUS channel_manager_setup_wlc_lcic(uvm_channel_pool_t *wlc_pool, uv
     return NV_OK;
 }
 
+static NV_STATUS channel_manager_create_conf_computing_pools(uvm_channel_manager_t *manager, unsigned *preferred_ce)
+{
+    NV_STATUS status;
+    unsigned wlc_lcic_ce_index;
+    uvm_channel_pool_t *sec2_pool = NULL;
+    uvm_channel_pool_t *wlc_pool = NULL;
+    uvm_channel_pool_t *lcic_pool = NULL;
+
+    if (!uvm_conf_computing_mode_enabled(manager->gpu))
+        return NV_OK;
+
+    status = uvm_rm_mem_alloc(manager->gpu,
+                             UVM_RM_MEM_TYPE_SYS,
+                             sizeof(UvmCslIv),
+                             UVM_CONF_COMPUTING_BUF_ALIGNMENT,
+                             &manager->gpu->conf_computing.iv_rm_mem);
+    if (status != NV_OK)
+        return status;
+
+    // Create SEC2 pool. This needs to be done first, initialization of
+    // other channels needs SEC2.
+    status = channel_pool_add(manager, UVM_CHANNEL_POOL_TYPE_SEC2, 0, &sec2_pool);
+    if (status != NV_OK)
+        return status;
+
+    manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_SEC2] = sec2_pool;
+
+    // Use the same CE as CPU TO GPU channels for WLC/LCIC
+    // Both need to use the same engine for the fixed schedule to work.
+    // TODO: Bug 3981928: [hcc][uvm] Optimize parameters of WLC/LCIC secure
+    // work launch
+    // Find a metric to select the best CE to use
+    wlc_lcic_ce_index = preferred_ce[UVM_CHANNEL_TYPE_CPU_TO_GPU];
+
+    // Create WLC/LCIC pools. This should be done early, CE channels use
+    // them for secure launch. The WLC pool must be created before the LCIC.
+    status = channel_pool_add(manager, UVM_CHANNEL_POOL_TYPE_WLC, wlc_lcic_ce_index, &wlc_pool);
+    if (status != NV_OK)
+        return status;
+
+    manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_WLC] = wlc_pool;
+
+    status = channel_pool_add(manager, UVM_CHANNEL_POOL_TYPE_LCIC, wlc_lcic_ce_index, &lcic_pool);
+    if (status != NV_OK)
+        return status;
+
+    status = channel_manager_setup_wlc_lcic(wlc_pool, lcic_pool);
+    if (status != NV_OK)
+        return status;
+
+    // The LCIC pool must be assigned after the call to
+    // channel_manager_setup_wlc_lcic(). It determines WLC and LCIC channels
+    // are ready to be used for secure work submission.
+    manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_LCIC] = lcic_pool;
+
+    return NV_OK;
+}
+
 static NV_STATUS channel_manager_create_pools(uvm_channel_manager_t *manager)
 {
     NV_STATUS status;
@@ -3154,62 +3173,11 @@ static NV_STATUS channel_manager_create_pools(uvm_channel_manager_t *manager)
     if (!manager->channel_pools)
         return NV_ERR_NO_MEMORY;
 
-    if (uvm_conf_computing_mode_enabled(manager->gpu)) {
-        uvm_channel_pool_t *sec2_pool = NULL;
-        uvm_channel_pool_t *wlc_pool = NULL;
-        uvm_channel_pool_t *lcic_pool = NULL;
-        unsigned wlc_lcic_ce_index;
-
-        status = uvm_rm_mem_alloc(manager->gpu,
-                                  UVM_RM_MEM_TYPE_SYS,
-                                  sizeof(UvmCslIv),
-                                  UVM_CONF_COMPUTING_BUF_ALIGNMENT,
-                                  &manager->gpu->conf_computing.iv_rm_mem);
-
-        if (status != NV_OK)
-            return status;
-
-        // Create SEC2 pool. This needs to be done first, initialization of
-        // other channels needs SEC2.
-        status = channel_pool_add_secure(manager, UVM_CHANNEL_POOL_TYPE_SEC2, 0, &sec2_pool);
-        if (status != NV_OK)
-            return status;
-
-        manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_SEC2] = sec2_pool;
-
-        // Use the same CE as CPU TO GPU channels for WLC/LCIC
-        // Both need to use the same engine for the fixed schedule to work.
-        // TODO: Bug 3981928: [hcc][uvm] Optimize parameters of WLC/LCIC secure
-        // work launch
-        // Find a metric to select the best CE to use
-        wlc_lcic_ce_index = preferred_ce[UVM_CHANNEL_TYPE_CPU_TO_GPU];
-
-        // Create WLC/LCIC pools. This should be done early, CE channels use
-        // them for secure launch. The WLC pool must be created before the LCIC.
-        status = channel_pool_add_secure(manager, UVM_CHANNEL_POOL_TYPE_WLC, wlc_lcic_ce_index, &wlc_pool);
-        if (status != NV_OK)
-            return status;
-
-        manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_WLC] = wlc_pool;
-
-        status = channel_pool_add_secure(manager, UVM_CHANNEL_POOL_TYPE_LCIC, wlc_lcic_ce_index, &lcic_pool);
-        if (status != NV_OK)
-            return status;
-
-        status = channel_manager_setup_wlc_lcic(wlc_pool, lcic_pool);
-        if (status != NV_OK)
-            return status;
-
-        // The LCIC pool must be assigned after the call to
-        // channel_manager_setup_wlc_lcic(). It determines WLC and LCIC channels
-        // are ready to be used for secure work submission.
-        manager->pool_to_use.default_for_type[UVM_CHANNEL_TYPE_LCIC] = lcic_pool;
-        status = channel_manager_create_ce_secure_pools(manager, preferred_ce);
-    }
-    else {
-        status = channel_manager_create_ce_pools(manager, preferred_ce);
-    }
+    status = channel_manager_create_conf_computing_pools(manager, preferred_ce);
+    if (status != NV_OK)
+        return status;
 
+    status = channel_manager_create_ce_pools(manager, preferred_ce);
     if (status != NV_OK)
         return status;
 
@@ -3292,7 +3260,17 @@ static void channel_manager_stop_wlc(uvm_channel_manager_t *manager)
 
         // Wait for the WLC/LCIC to be primed. This means that PUT == GET + 2
         // and a WLC doorbell ring is enough to start work.
-        UVM_SPIN_WHILE(!uvm_gpu_tracking_semaphore_is_completed(&channel->tracking_sem), &spin);
+        UVM_SPIN_WHILE(!uvm_gpu_tracking_semaphore_is_completed(&channel->tracking_sem), &spin) {
+            status = uvm_channel_check_errors(channel);
+            if (status != NV_OK) {
+                UVM_ERR_PRINT_NV_STATUS("Failed to wait for LCIC channel (%s) completion", status, channel->name);
+                break;
+            }
+        }
+
+        // Continue on error and attempt to stop WLC below. This can lead to
+        // channel destruction with mismatched GET and PUT pointers. RM will
+        // print errors if that's the case, but channel destruction succeeeds.
     }
 
     status = uvm_push_begin(manager, UVM_CHANNEL_TYPE_SEC2, &push, "Stop WLC channels");

kernel-open/nvidia-uvm/uvm_channel.h

@@ -104,16 +104,14 @@ typedef enum
     // ----------------------------------
     // Channel type with fixed schedules
 
-    // Work Launch Channel (WLC) is a specialized channel
-    // for launching work on other channels when
-    // Confidential Computing is enabled.
-    // It is paired with LCIC (below)
+    // Work Launch Channel (WLC) is a specialized channel for launching work on
+    // other channels when the Confidential Computing is feature enabled. It is
+    // paired with LCIC (below)
     UVM_CHANNEL_TYPE_WLC,
 
-    // Launch Confirmation Indicator Channel (LCIC) is a
-    // specialized channel with fixed schedule. It gets
-    // triggered by executing WLC work, and makes sure that
-    // WLC get/put pointers are up-to-date.
+    // Launch Confirmation Indicator Channel (LCIC) is a specialized channel
+    // with fixed schedule. It gets triggered by executing WLC work, and makes
+    // sure that WLC get/put pointers are up-to-date.
     UVM_CHANNEL_TYPE_LCIC,
 
     UVM_CHANNEL_TYPE_COUNT,
@@ -242,11 +240,9 @@ typedef struct
     DECLARE_BITMAP(push_locks, UVM_CHANNEL_MAX_NUM_CHANNELS_PER_POOL);
 
     // Counting semaphore for available and unlocked channels, it must be
-    // acquired before submitting work to a secure channel.
+    // acquired before submitting work to a channel when the Confidential
+    // Computing feature is enabled.
     uvm_semaphore_t push_sem;
-
-    // See uvm_channel_is_secure() documentation.
-    bool secure;
 } uvm_channel_pool_t;
 
 struct uvm_channel_struct
@@ -304,8 +300,9 @@ struct uvm_channel_struct
         // its internal operation and each push may modify this state.
         uvm_mutex_t push_lock;
 
-        // Every secure channel has cryptographic state in HW, which is
-        // mirrored here for CPU-side operations.
+        // When the Confidential Computing feature is enabled, every channel has
+        // cryptographic state in HW, which is mirrored here for CPU-side
+        // operations.
         UvmCslContext ctx;
         bool is_ctx_initialized;
 
@@ -459,46 +456,28 @@ struct uvm_channel_manager_struct
 // Create a channel manager for the GPU
 NV_STATUS uvm_channel_manager_create(uvm_gpu_t *gpu, uvm_channel_manager_t **manager_out);
 
-static bool uvm_channel_pool_is_ce(uvm_channel_pool_t *pool);
-
-// A channel is secure if it has HW encryption capabilities.
-//
-// Secure channels are treated differently in the UVM driver. Each secure
-// channel has a unique CSL context associated with it, has relatively
-// restrictive reservation policies (in comparison with non-secure channels),
-// it is requested to be allocated differently by RM, etc.
-static bool uvm_channel_pool_is_secure(uvm_channel_pool_t *pool)
+static bool uvm_pool_type_is_valid(uvm_channel_pool_type_t pool_type)
 {
-    return pool->secure;
-}
-
-static bool uvm_channel_is_secure(uvm_channel_t *channel)
-{
-    return uvm_channel_pool_is_secure(channel->pool);
+    return (is_power_of_2(pool_type) && (pool_type < UVM_CHANNEL_POOL_TYPE_MASK));
 }
 
 static bool uvm_channel_pool_is_sec2(uvm_channel_pool_t *pool)
 {
-    UVM_ASSERT(pool->pool_type < UVM_CHANNEL_POOL_TYPE_MASK);
+    UVM_ASSERT(uvm_pool_type_is_valid(pool->pool_type));
 
     return (pool->pool_type == UVM_CHANNEL_POOL_TYPE_SEC2);
 }
 
-static bool uvm_channel_pool_is_secure_ce(uvm_channel_pool_t *pool)
-{
-    return uvm_channel_pool_is_secure(pool) && uvm_channel_pool_is_ce(pool);
-}
-
 static bool uvm_channel_pool_is_wlc(uvm_channel_pool_t *pool)
 {
-    UVM_ASSERT(pool->pool_type < UVM_CHANNEL_POOL_TYPE_MASK);
+    UVM_ASSERT(uvm_pool_type_is_valid(pool->pool_type));
 
     return (pool->pool_type == UVM_CHANNEL_POOL_TYPE_WLC);
 }
 
 static bool uvm_channel_pool_is_lcic(uvm_channel_pool_t *pool)
 {
-    UVM_ASSERT(pool->pool_type < UVM_CHANNEL_POOL_TYPE_MASK);
+    UVM_ASSERT(uvm_pool_type_is_valid(pool->pool_type));
 
     return (pool->pool_type == UVM_CHANNEL_POOL_TYPE_LCIC);
 }
@@ -508,11 +487,6 @@ static bool uvm_channel_is_sec2(uvm_channel_t *channel)
     return uvm_channel_pool_is_sec2(channel->pool);
 }
 
-static bool uvm_channel_is_secure_ce(uvm_channel_t *channel)
-{
-    return uvm_channel_pool_is_secure_ce(channel->pool);
-}
-
 static bool uvm_channel_is_wlc(uvm_channel_t *channel)
 {
     return uvm_channel_pool_is_wlc(channel->pool);
@@ -523,12 +497,9 @@ static bool uvm_channel_is_lcic(uvm_channel_t *channel)
     return uvm_channel_pool_is_lcic(channel->pool);
 }
 
-bool uvm_channel_type_requires_secure_pool(uvm_gpu_t *gpu, uvm_channel_type_t channel_type);
-NV_STATUS uvm_channel_secure_init(uvm_gpu_t *gpu, uvm_channel_t *channel);
-
 static bool uvm_channel_pool_is_proxy(uvm_channel_pool_t *pool)
 {
-    UVM_ASSERT(pool->pool_type < UVM_CHANNEL_POOL_TYPE_MASK);
+    UVM_ASSERT(uvm_pool_type_is_valid(pool->pool_type));
 
     return pool->pool_type == UVM_CHANNEL_POOL_TYPE_CE_PROXY;
 }
@@ -540,11 +511,7 @@ static bool uvm_channel_is_proxy(uvm_channel_t *channel)
 
 static bool uvm_channel_pool_is_ce(uvm_channel_pool_t *pool)
 {
-    UVM_ASSERT(pool->pool_type < UVM_CHANNEL_POOL_TYPE_MASK);
-    if (uvm_channel_pool_is_wlc(pool) || uvm_channel_pool_is_lcic(pool))
-        return true;
-
-    return (pool->pool_type == UVM_CHANNEL_POOL_TYPE_CE) || uvm_channel_pool_is_proxy(pool);
+    return !uvm_channel_pool_is_sec2(pool);
 }
 
 static bool uvm_channel_is_ce(uvm_channel_t *channel)
@@ -686,6 +653,11 @@ static uvm_gpu_t *uvm_channel_get_gpu(uvm_channel_t *channel)
     return channel->pool->manager->gpu;
 }
 
+static uvm_pushbuffer_t *uvm_channel_get_pushbuffer(uvm_channel_t *channel)
+{
+    return channel->pool->manager->pushbuffer;
+}
+
 // Index of a channel within the owning pool
 static unsigned uvm_channel_index_in_pool(const uvm_channel_t *channel)
 {

kernel-open/nvidia-uvm/uvm_channel_test.c

@@ -681,9 +681,10 @@ done:
 }
 
 // The following test is inspired by uvm_push_test.c:test_concurrent_pushes.
-// This test verifies that concurrent pushes using the same secure channel pool
-// select different channels.
-NV_STATUS test_secure_channel_selection(uvm_va_space_t *va_space)
+// This test verifies that concurrent pushes using the same channel pool
+// select different channels, when the Confidential Computing feature is
+// enabled.
+NV_STATUS test_conf_computing_channel_selection(uvm_va_space_t *va_space)
 {
     NV_STATUS status = NV_OK;
     uvm_channel_pool_t *pool;
@@ -703,9 +704,6 @@ NV_STATUS test_secure_channel_selection(uvm_va_space_t *va_space)
         uvm_channel_type_t channel_type;
 
         for (channel_type = 0; channel_type < UVM_CHANNEL_TYPE_COUNT; channel_type++) {
-            if (!uvm_channel_type_requires_secure_pool(gpu, channel_type))
-                continue;
-
             pool = gpu->channel_manager->pool_to_use.default_for_type[channel_type];
             TEST_CHECK_RET(pool != NULL);
 
@@ -997,7 +995,7 @@ NV_STATUS uvm_test_channel_sanity(UVM_TEST_CHANNEL_SANITY_PARAMS *params, struct
     if (status != NV_OK)
         goto done;
 
-    status = test_secure_channel_selection(va_space);
+    status = test_conf_computing_channel_selection(va_space);
     if (status != NV_OK)
         goto done;
 

kernel-open/nvidia-uvm/uvm_common.h

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2013-2021 NVIDIA Corporation
+    Copyright (c) 2013-2024 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -21,8 +21,8 @@
 
 *******************************************************************************/
 
-#ifndef _UVM_COMMON_H
-#define _UVM_COMMON_H
+#ifndef __UVM_COMMON_H__
+#define __UVM_COMMON_H__
 
 #ifdef DEBUG
     #define UVM_IS_DEBUG() 1
@@ -413,4 +413,42 @@ static inline void uvm_touch_page(struct page *page)
 // Return true if the VMA is one used by UVM managed allocations.
 bool uvm_vma_is_managed(struct vm_area_struct *vma);
 
-#endif /* _UVM_COMMON_H */
+static bool uvm_platform_uses_canonical_form_address(void)
+{
+    if (NVCPU_IS_PPC64LE)
+        return false;
+
+    return true;
+}
+
+// Similar to the GPU MMU HAL num_va_bits(), it returns the CPU's num_va_bits().
+static NvU32 uvm_cpu_num_va_bits(void)
+{
+    return fls64(TASK_SIZE - 1) + 1;
+}
+
+// Return the unaddressable range in a num_va_bits-wide VA space, [first, outer)
+static void uvm_get_unaddressable_range(NvU32 num_va_bits, NvU64 *first, NvU64 *outer)
+{
+    UVM_ASSERT(num_va_bits < 64);
+    UVM_ASSERT(first);
+    UVM_ASSERT(outer);
+
+    // Maxwell GPUs (num_va_bits == 40b) do not support canonical form address
+    // even when plugged into platforms using it.
+    if (uvm_platform_uses_canonical_form_address() && num_va_bits > 40) {
+        *first = 1ULL << (num_va_bits - 1);
+        *outer = (NvU64)((NvS64)(1ULL << 63) >> (64 - num_va_bits));
+    }
+    else {
+        *first = 1ULL << num_va_bits;
+        *outer = ~0Ull;
+    }
+}
+
+static void uvm_cpu_get_unaddressable_range(NvU64 *first, NvU64 *outer)
+{
+    return uvm_get_unaddressable_range(uvm_cpu_num_va_bits(), first, outer);
+}
+
+#endif /* __UVM_COMMON_H__ */

kernel-open/nvidia-uvm/uvm_conf_computing.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2021 NVIDIA Corporation
+    Copyright (c) 2021-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -54,23 +54,26 @@ bool uvm_conf_computing_mode_is_hcc(const uvm_gpu_t *gpu)
     return uvm_conf_computing_get_mode(gpu->parent) == UVM_GPU_CONF_COMPUTE_MODE_HCC;
 }
 
-NV_STATUS uvm_conf_computing_init_parent_gpu(const uvm_parent_gpu_t *parent)
+void uvm_conf_computing_check_parent_gpu(const uvm_parent_gpu_t *parent)
 {
-    UvmGpuConfComputeMode cc, sys_cc;
-    uvm_gpu_t *first;
+    uvm_gpu_t *first_gpu;
 
     uvm_assert_mutex_locked(&g_uvm_global.global_lock);
 
+    // The Confidential Computing state of the GPU should match that of the
+    // system.
+    UVM_ASSERT(uvm_conf_computing_mode_enabled_parent(parent) == g_uvm_global.conf_computing_enabled);
+
     // TODO: Bug 2844714: since we have no routine to traverse parent GPUs,
     // find first child GPU and get its parent.
-    first = uvm_global_processor_mask_find_first_gpu(&g_uvm_global.retained_gpus);
-    if (!first)
-        return NV_OK;
+    first_gpu = uvm_global_processor_mask_find_first_gpu(&g_uvm_global.retained_gpus);
+    if (first_gpu == NULL)
+        return;
 
-    sys_cc = uvm_conf_computing_get_mode(first->parent);
-    cc = uvm_conf_computing_get_mode(parent);
-
-    return cc == sys_cc ? NV_OK : NV_ERR_NOT_SUPPORTED;
+    // All GPUs derive Confidential Computing status from their parent. By
+    // current policy all parent GPUs have identical Confidential Computing
+    // status.
+    UVM_ASSERT(uvm_conf_computing_get_mode(parent) == uvm_conf_computing_get_mode(first_gpu->parent));
 }
 
 static void dma_buffer_destroy_locked(uvm_conf_computing_dma_buffer_pool_t *dma_buffer_pool,

kernel-open/nvidia-uvm/uvm_conf_computing.h

@@ -60,10 +60,8 @@
 // UVM_METHOD_SIZE * 2 * 10 = 80.
 #define UVM_CONF_COMPUTING_SIGN_BUF_MAX_SIZE 80
 
-// All GPUs derive confidential computing status from their parent.
-// By current policy all parent GPUs have identical confidential
-// computing status.
-NV_STATUS uvm_conf_computing_init_parent_gpu(const uvm_parent_gpu_t *parent);
+void uvm_conf_computing_check_parent_gpu(const uvm_parent_gpu_t *parent);
+
 bool uvm_conf_computing_mode_enabled_parent(const uvm_parent_gpu_t *parent);
 bool uvm_conf_computing_mode_enabled(const uvm_gpu_t *gpu);
 bool uvm_conf_computing_mode_is_hcc(const uvm_gpu_t *gpu);

kernel-open/nvidia-uvm/uvm_global.c

@@ -71,11 +71,6 @@ static void uvm_unregister_callbacks(void)
     }
 }
 
-static void sev_init(const UvmPlatformInfo *platform_info)
-{
-    g_uvm_global.sev_enabled = platform_info->sevEnabled;
-}
-
 NV_STATUS uvm_global_init(void)
 {
     NV_STATUS status;
@@ -124,8 +119,7 @@ NV_STATUS uvm_global_init(void)
 
     uvm_ats_init(&platform_info);
     g_uvm_global.num_simulated_devices = 0;
-
-    sev_init(&platform_info);
+    g_uvm_global.conf_computing_enabled = platform_info.confComputingEnabled;
 
     status = uvm_gpu_init();
     if (status != NV_OK) {

kernel-open/nvidia-uvm/uvm_global.h

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2021 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -143,11 +143,16 @@ struct uvm_global_struct
         struct page *page;
     } unload_state;
 
-    // AMD Secure Encrypted Virtualization (SEV) status. True if VM has SEV
-    // enabled. This field is set once during global initialization
-    // (uvm_global_init), and can be read afterwards without acquiring any
-    // locks.
-    bool sev_enabled;
+    // True if the VM has AMD's SEV, or equivalent HW security extensions such
+    // as Intel's TDX, enabled. The flag is always false on the host.
+    //
+    // This value moves in tandem with that of Confidential Computing in the
+    // GPU(s) in all supported configurations, so it is used as a proxy for the
+    // Confidential Computing state.
+    //
+    // This field is set once during global initialization (uvm_global_init),
+    // and can be read afterwards without acquiring any locks.
+    bool conf_computing_enabled;
 };
 
 // Initialize global uvm state

kernel-open/nvidia-uvm/uvm_gpu.c

@@ -218,19 +218,12 @@ static bool gpu_supports_uvm(uvm_parent_gpu_t *parent_gpu)
     return parent_gpu->rm_info.subdeviceCount == 1;
 }
 
-static bool platform_uses_canonical_form_address(void)
-{
-    if (NVCPU_IS_PPC64LE)
-        return false;
-
-    return true;
-}
-
 bool uvm_gpu_can_address(uvm_gpu_t *gpu, NvU64 addr, NvU64 size)
 {
     // Lower and upper address spaces are typically found in platforms that use
     // the canonical address form.
     NvU64 max_va_lower;
+    NvU64 min_va_upper;
     NvU64 addr_end = addr + size - 1;
     NvU8 gpu_addr_shift;
     NvU8 cpu_addr_shift;
@@ -243,7 +236,7 @@ bool uvm_gpu_can_address(uvm_gpu_t *gpu, NvU64 addr, NvU64 size)
     UVM_ASSERT(size > 0);
 
     gpu_addr_shift = gpu->address_space_tree.hal->num_va_bits();
-    cpu_addr_shift = fls64(TASK_SIZE - 1) + 1;
+    cpu_addr_shift = uvm_cpu_num_va_bits();
     addr_shift = gpu_addr_shift;
 
     // Pascal+ GPUs are capable of accessing kernel pointers in various modes
@@ -279,9 +272,7 @@ bool uvm_gpu_can_address(uvm_gpu_t *gpu, NvU64 addr, NvU64 size)
     //               0 +----------------+               0 +----------------+
 
     // On canonical form address platforms and Pascal+ GPUs.
-    if (platform_uses_canonical_form_address() && gpu_addr_shift > 40) {
-        NvU64 min_va_upper;
-
+    if (uvm_platform_uses_canonical_form_address() && gpu_addr_shift > 40) {
         // On x86, when cpu_addr_shift > gpu_addr_shift, it means the CPU uses
         // 5-level paging and the GPU is pre-Hopper. On Pascal-Ada GPUs (49b
         // wide VA) we set addr_shift to match a 4-level paging x86 (48b wide).
@@ -292,15 +283,11 @@ bool uvm_gpu_can_address(uvm_gpu_t *gpu, NvU64 addr, NvU64 size)
             addr_shift = gpu_addr_shift;
         else
             addr_shift = cpu_addr_shift;
+    }
 
-        min_va_upper = (NvU64)((NvS64)(1ULL << 63) >> (64 - addr_shift));
-        max_va_lower = 1ULL << (addr_shift - 1);
-        return (addr_end < max_va_lower) || (addr >= min_va_upper);
-    }
-    else {
-        max_va_lower = 1ULL << addr_shift;
-        return addr_end < max_va_lower;
-    }
+    uvm_get_unaddressable_range(addr_shift, &max_va_lower, &min_va_upper);
+
+    return (addr_end < max_va_lower) || (addr >= min_va_upper);
 }
 
 // The internal UVM VAS does not use canonical form addresses.
@@ -326,14 +313,14 @@ NvU64 uvm_parent_gpu_canonical_address(uvm_parent_gpu_t *parent_gpu, NvU64 addr)
     NvU8 addr_shift;
     NvU64 input_addr = addr;
 
-    if (platform_uses_canonical_form_address()) {
+    if (uvm_platform_uses_canonical_form_address()) {
         // When the CPU VA width is larger than GPU's, it means that:
         // On ARM: the CPU is on LVA mode and the GPU is pre-Hopper.
         // On x86: the CPU uses 5-level paging and the GPU is pre-Hopper.
         // We sign-extend on the 48b on ARM and on the 47b on x86 to mirror the
         // behavior of CPUs with smaller (than GPU) VA widths.
         gpu_addr_shift = parent_gpu->arch_hal->mmu_mode_hal(UVM_PAGE_SIZE_64K)->num_va_bits();
-        cpu_addr_shift = fls64(TASK_SIZE - 1) + 1;
+        cpu_addr_shift = uvm_cpu_num_va_bits();
 
         if (cpu_addr_shift > gpu_addr_shift)
             addr_shift = NVCPU_IS_X86_64 ? 48 : 49;
@@ -1099,12 +1086,7 @@ static NV_STATUS init_parent_gpu(uvm_parent_gpu_t *parent_gpu,
         return status;
     }
 
-    status = uvm_conf_computing_init_parent_gpu(parent_gpu);
-    if (status != NV_OK) {
-        UVM_ERR_PRINT("Confidential computing: %s, GPU %s\n",
-                      nvstatusToString(status), parent_gpu->name);
-        return status;
-    }
+    uvm_conf_computing_check_parent_gpu(parent_gpu);
 
     parent_gpu->pci_dev = gpu_platform_info->pci_dev;
     parent_gpu->closest_cpu_numa_node = dev_to_node(&parent_gpu->pci_dev->dev);

kernel-open/nvidia-uvm/uvm_gpu.h

@@ -46,6 +46,7 @@
 #include "uvm_rb_tree.h"
 #include "uvm_perf_prefetch.h"
 #include "nv-kthread-q.h"
+#include <linux/mmu_notifier.h>
 #include "uvm_conf_computing.h"
 
 // Buffer length to store uvm gpu id, RM device name and gpu uuid.
@@ -56,14 +57,16 @@
 
 typedef struct
 {
-    // Number of faults from this uTLB that have been fetched but have not been serviced yet
+    // Number of faults from this uTLB that have been fetched but have not been
+    // serviced yet.
     NvU32 num_pending_faults;
 
     // Whether the uTLB contains fatal faults
     bool has_fatal_faults;
 
-    // We have issued a replay of type START_ACK_ALL while containing fatal faults. This puts
-    // the uTLB in lockdown mode and no new translations are accepted
+    // We have issued a replay of type START_ACK_ALL while containing fatal
+    // faults. This puts the uTLB in lockdown mode and no new translations are
+    // accepted.
     bool in_lockdown;
 
     // We have issued a cancel on this uTLB
@@ -125,8 +128,8 @@ struct uvm_service_block_context_struct
         struct list_head service_context_list;
 
         // A mask of GPUs that need to be checked for ECC errors before the CPU
-        // fault handler returns, but after the VA space lock has been unlocked to
-        // avoid the RM/UVM VA space lock deadlocks.
+        // fault handler returns, but after the VA space lock has been unlocked
+        // to avoid the RM/UVM VA space lock deadlocks.
         uvm_processor_mask_t gpus_to_check_for_ecc;
 
         // This is set to throttle page fault thrashing.
@@ -159,9 +162,9 @@ struct uvm_service_block_context_struct
 
     struct
     {
-        // Per-processor mask with the pages that will be resident after servicing.
-        // We need one mask per processor because we may coalesce faults that
-        // trigger migrations to different processors.
+        // Per-processor mask with the pages that will be resident after
+        // servicing. We need one mask per processor because we may coalesce
+        // faults that trigger migrations to different processors.
         uvm_page_mask_t new_residency;
     } per_processor_masks[UVM_ID_MAX_PROCESSORS];
 
@@ -178,23 +181,28 @@ struct uvm_service_block_context_struct
 typedef struct
 {
     // Mask of read faulted pages in a UVM_VA_BLOCK_SIZE aligned region of a SAM
-    // VMA. Used for batching ATS faults in a vma.
+    // VMA. Used for batching ATS faults in a vma. This is unused for access
+    // counter service requests.
     uvm_page_mask_t read_fault_mask;
 
     // Mask of write faulted pages in a UVM_VA_BLOCK_SIZE aligned region of a
-    // SAM VMA. Used for batching ATS faults in a vma.
+    // SAM VMA. Used for batching ATS faults in a vma. This is unused for access
+    // counter service requests.
     uvm_page_mask_t write_fault_mask;
 
     // Mask of successfully serviced pages in a UVM_VA_BLOCK_SIZE aligned region
-    // of a SAM VMA. Used to return ATS fault status.
+    // of a SAM VMA. Used to return ATS fault status. This is unused for access
+    // counter service requests.
     uvm_page_mask_t faults_serviced_mask;
 
     // Mask of successfully serviced read faults on pages in write_fault_mask.
+    // This is unused for access counter service requests.
     uvm_page_mask_t reads_serviced_mask;
 
-    // Temporary mask used for uvm_page_mask_or_equal. This is used since
-    // bitmap_or_equal() isn't present in all linux kernel versions.
-    uvm_page_mask_t tmp_mask;
+    // Mask of all accessed pages in a UVM_VA_BLOCK_SIZE aligned region of a SAM
+    // VMA. This is used as input for access counter service requests and output
+    // of fault service requests.
+    uvm_page_mask_t accessed_mask;
 
     // Client type of the service requestor.
     uvm_fault_client_type_t client_type;
@@ -204,6 +212,40 @@ typedef struct
 
     // New residency NUMA node ID of the faulting region.
     int residency_node;
+
+    struct
+    {
+        // True if preferred_location was set on this faulting region.
+        // UVM_VA_BLOCK_SIZE sized region in the faulting region bound by the
+        // VMA is is prefetched if preferred_location was set and if first_touch
+        // is true;
+        bool has_preferred_location;
+
+        // True if the UVM_VA_BLOCK_SIZE sized region isn't resident on any
+        // node. False if any page in the region is resident somewhere.
+        bool first_touch;
+
+        // Mask of prefetched pages in a UVM_VA_BLOCK_SIZE aligned region of a
+        // SAM VMA.
+        uvm_page_mask_t prefetch_pages_mask;
+
+        // PFN info of the faulting region
+        unsigned long pfns[PAGES_PER_UVM_VA_BLOCK];
+
+        // Faulting/preferred processor residency mask of the faulting region.
+        uvm_page_mask_t residency_mask;
+
+#if defined(NV_MMU_INTERVAL_NOTIFIER)
+        // MMU notifier used to compute residency of this faulting region.
+        struct mmu_interval_notifier notifier;
+#endif
+
+        uvm_va_space_t *va_space;
+
+        // Prefetch temporary state.
+        uvm_perf_prefetch_bitmap_tree_t bitmap_tree;
+    } prefetch_state;
+
 } uvm_ats_fault_context_t;
 
 struct uvm_fault_service_batch_context_struct
@@ -228,7 +270,10 @@ struct uvm_fault_service_batch_context_struct
 
     NvU32 num_coalesced_faults;
 
-    bool has_fatal_faults;
+    // One of the VA spaces in this batch which had fatal faults. If NULL, no
+    // faults were fatal. More than one VA space could have fatal faults, but we
+    // pick one to be the target of the cancel sequence.
+    uvm_va_space_t *fatal_va_space;
 
     bool has_throttled_faults;
 
@@ -256,11 +301,8 @@ struct uvm_fault_service_batch_context_struct
 
 struct uvm_ats_fault_invalidate_struct
 {
-    // Whether the TLB batch contains any information
-    bool            write_faults_in_batch;
-
-    // Batch of TLB entries to be invalidated
-    uvm_tlb_batch_t write_faults_tlb_batch;
+    bool            tlb_batch_pending;
+    uvm_tlb_batch_t tlb_batch;
 };
 
 typedef struct
@@ -405,20 +447,9 @@ struct uvm_access_counter_service_batch_context_struct
         NvU32                             num_notifications;
 
         // Boolean used to avoid sorting the fault batch by instance_ptr if we
-        // determine at fetch time that all the access counter notifications in the
-        // batch report the same instance_ptr
+        // determine at fetch time that all the access counter notifications in
+        // the batch report the same instance_ptr
         bool is_single_instance_ptr;
-
-        // Scratch space, used to generate artificial physically addressed notifications.
-        // Virtual address notifications are always aligned to 64k. This means up to 16
-        // different physical locations could have been accessed to trigger one notification.
-        // The sub-granularity mask can correspond to any of them.
-        struct
-        {
-            uvm_processor_id_t resident_processors[16];
-            uvm_gpu_phys_address_t phys_addresses[16];
-            uvm_access_counter_buffer_entry_t phys_entry;
-        } scratch;
     } virt;
 
     struct
@@ -429,8 +460,8 @@ struct uvm_access_counter_service_batch_context_struct
         NvU32                              num_notifications;
 
         // Boolean used to avoid sorting the fault batch by aperture if we
-        // determine at fetch time that all the access counter notifications in the
-        // batch report the same aperture
+        // determine at fetch time that all the access counter notifications in
+        // the batch report the same aperture
         bool                              is_single_aperture;
     } phys;
 
@@ -440,6 +471,9 @@ struct uvm_access_counter_service_batch_context_struct
     // Structure used to coalesce access counter servicing in a VA block
     uvm_service_block_context_t block_service_context;
 
+    // Structure used to service access counter migrations in an ATS block.
+    uvm_ats_fault_context_t ats_context;
+
     // Unique id (per-GPU) generated for tools events recording
     NvU32 batch_id;
 };
@@ -626,8 +660,8 @@ struct uvm_gpu_struct
     struct
     {
         // Big page size used by the internal UVM VA space
-        // Notably it may be different than the big page size used by a user's VA
-        // space in general.
+        // Notably it may be different than the big page size used by a user's
+        // VA space in general.
         NvU32 internal_size;
     } big_page;
 
@@ -653,8 +687,8 @@ struct uvm_gpu_struct
         // lazily-populated array of peer GPUs, indexed by the peer's GPU index
         uvm_gpu_t *peer_gpus[UVM_ID_MAX_GPUS];
 
-        // Leaf spinlock used to synchronize access to the peer_gpus table so that
-        // it can be safely accessed from the access counters bottom half
+        // Leaf spinlock used to synchronize access to the peer_gpus table so
+        // that it can be safely accessed from the access counters bottom half
         uvm_spinlock_t peer_gpus_lock;
     } peer_info;
 
@@ -945,6 +979,10 @@ struct uvm_parent_gpu_struct
 
     bool plc_supported;
 
+    // If true, page_tree initialization pre-populates no_ats_ranges. It only
+    // affects ATS systems.
+    bool no_ats_range_required;
+
     // Parameters used by the TLB batching API
     struct
     {
@@ -1016,14 +1054,16 @@ struct uvm_parent_gpu_struct
     // Interrupt handling state and locks
     uvm_isr_info_t isr;
 
-    // Fault buffer info. This is only valid if supports_replayable_faults is set to true
+    // Fault buffer info. This is only valid if supports_replayable_faults is
+    // set to true.
     uvm_fault_buffer_info_t fault_buffer_info;
 
     // PMM lazy free processing queue.
     // TODO: Bug 3881835: revisit whether to use nv_kthread_q_t or workqueue.
     nv_kthread_q_t lazy_free_q;
 
-    // Access counter buffer info. This is only valid if supports_access_counters is set to true
+    // Access counter buffer info. This is only valid if
+    // supports_access_counters is set to true.
     uvm_access_counter_buffer_info_t access_counter_buffer_info;
 
     // Number of uTLBs per GPC. This information is only valid on Pascal+ GPUs.
@@ -1073,7 +1113,7 @@ struct uvm_parent_gpu_struct
     uvm_rb_tree_t instance_ptr_table;
     uvm_spinlock_t instance_ptr_table_lock;
 
-    // This is set to true if the GPU belongs to an SLI group. Else, set to false.
+    // This is set to true if the GPU belongs to an SLI group.
     bool sli_enabled;
 
     struct
@@ -1100,8 +1140,8 @@ struct uvm_parent_gpu_struct
     // environment, rather than using the peer-id field of the PTE (which can
     // only address 8 gpus), all gpus are assigned a 47-bit physical address
     // space by the fabric manager. Any physical address access to these
-    // physical address spaces are routed through the switch to the corresponding
-    // peer.
+    // physical address spaces are routed through the switch to the
+    // corresponding peer.
     struct
     {
         bool is_nvswitch_connected;
@@ -1127,6 +1167,16 @@ struct uvm_parent_gpu_struct
         NvU64 memory_window_start;
         NvU64 memory_window_end;
     } system_bus;
+
+    // WAR to issue ATS TLB invalidation commands ourselves.
+    struct
+    {
+        uvm_mutex_t smmu_lock;
+        struct page *smmu_cmdq;
+        void __iomem *smmu_cmdqv_base;
+        unsigned long smmu_prod;
+        unsigned long smmu_cons;
+    } smmu_war;
 };
 
 static const char *uvm_gpu_name(uvm_gpu_t *gpu)
@@ -1316,7 +1366,8 @@ void uvm_gpu_release_pcie_peer_access(uvm_gpu_t *gpu0, uvm_gpu_t *gpu1);
 // They must not be the same gpu.
 uvm_aperture_t uvm_gpu_peer_aperture(uvm_gpu_t *local_gpu, uvm_gpu_t *remote_gpu);
 
-// Get the processor id accessible by the given GPU for the given physical address
+// Get the processor id accessible by the given GPU for the given physical
+// address.
 uvm_processor_id_t uvm_gpu_get_processor_id_by_address(uvm_gpu_t *gpu, uvm_gpu_phys_address_t addr);
 
 // Get the P2P capabilities between the gpus with the given indexes
@@ -1413,9 +1464,9 @@ NV_STATUS uvm_gpu_check_ecc_error(uvm_gpu_t *gpu);
 
 // Check for ECC errors without calling into RM
 //
-// Calling into RM is problematic in many places, this check is always safe to do.
-// Returns NV_WARN_MORE_PROCESSING_REQUIRED if there might be an ECC error and
-// it's required to call uvm_gpu_check_ecc_error() to be sure.
+// Calling into RM is problematic in many places, this check is always safe to
+// do. Returns NV_WARN_MORE_PROCESSING_REQUIRED if there might be an ECC error
+// and it's required to call uvm_gpu_check_ecc_error() to be sure.
 NV_STATUS uvm_gpu_check_ecc_error_no_rm(uvm_gpu_t *gpu);
 
 // Map size bytes of contiguous sysmem on the GPU for physical access
@@ -1472,6 +1523,8 @@ bool uvm_gpu_can_address(uvm_gpu_t *gpu, NvU64 addr, NvU64 size);
 // The GPU must be initialized before calling this function.
 bool uvm_gpu_can_address_kernel(uvm_gpu_t *gpu, NvU64 addr, NvU64 size);
 
+bool uvm_platform_uses_canonical_form_address(void);
+
 // Returns addr's canonical form for host systems that use canonical form
 // addresses.
 NvU64 uvm_parent_gpu_canonical_address(uvm_parent_gpu_t *parent_gpu, NvU64 addr);
@@ -1518,8 +1571,9 @@ uvm_aperture_t uvm_gpu_page_tree_init_location(const uvm_gpu_t *gpu);
 // Debug print of GPU properties
 void uvm_gpu_print(uvm_gpu_t *gpu);
 
-// Add the given instance pointer -> user_channel mapping to this GPU. The bottom
-// half GPU page fault handler uses this to look up the VA space for GPU faults.
+// Add the given instance pointer -> user_channel mapping to this GPU. The
+// bottom half GPU page fault handler uses this to look up the VA space for GPU
+// faults.
 NV_STATUS uvm_gpu_add_user_channel(uvm_gpu_t *gpu, uvm_user_channel_t *user_channel);
 void uvm_gpu_remove_user_channel(uvm_gpu_t *gpu, uvm_user_channel_t *user_channel);
 

kernel-open/nvidia-uvm/uvm_gpu_access_counters.c

@@ -33,17 +33,18 @@
 #include "uvm_va_space_mm.h"
 #include "uvm_pmm_sysmem.h"
 #include "uvm_perf_module.h"
+#include "uvm_ats.h"
+#include "uvm_ats_faults.h"
 
 #define UVM_PERF_ACCESS_COUNTER_BATCH_COUNT_MIN     1
 #define UVM_PERF_ACCESS_COUNTER_BATCH_COUNT_DEFAULT 256
-#define UVM_PERF_ACCESS_COUNTER_GRANULARITY_DEFAULT "2m"
+#define UVM_PERF_ACCESS_COUNTER_GRANULARITY         UVM_ACCESS_COUNTER_GRANULARITY_2M
 #define UVM_PERF_ACCESS_COUNTER_THRESHOLD_MIN       1
 #define UVM_PERF_ACCESS_COUNTER_THRESHOLD_MAX       ((1 << 16) - 1)
 #define UVM_PERF_ACCESS_COUNTER_THRESHOLD_DEFAULT   256
 
-#define UVM_ACCESS_COUNTER_ACTION_NOTIFY 0x1
-#define UVM_ACCESS_COUNTER_ACTION_CLEAR  0x2
-#define UVM_ACCESS_COUNTER_ON_MANAGED    0x4
+#define UVM_ACCESS_COUNTER_ACTION_CLEAR     0x1
+#define UVM_ACCESS_COUNTER_PHYS_ON_MANAGED  0x2
 
 // Each page in a tracked physical range may belong to a different VA Block. We
 // preallocate an array of reverse map translations. However, access counter
@@ -54,12 +55,6 @@
 #define UVM_MAX_TRANSLATION_SIZE (2 * 1024 * 1024ULL)
 #define UVM_SUB_GRANULARITY_REGIONS 32
 
-// The GPU offers the following tracking granularities: 64K, 2M, 16M, 16G
-//
-// Use the largest granularity to minimize the number of access counter
-// notifications. This is fine because we simply drop the notifications during
-// normal operation, and tests override these values.
-static UVM_ACCESS_COUNTER_GRANULARITY g_uvm_access_counter_granularity;
 static unsigned g_uvm_access_counter_threshold;
 
 // Per-VA space access counters information
@@ -87,7 +82,6 @@ static int uvm_perf_access_counter_momc_migration_enable = -1;
 static unsigned uvm_perf_access_counter_batch_count = UVM_PERF_ACCESS_COUNTER_BATCH_COUNT_DEFAULT;
 
 // See module param documentation below
-static char *uvm_perf_access_counter_granularity = UVM_PERF_ACCESS_COUNTER_GRANULARITY_DEFAULT;
 static unsigned uvm_perf_access_counter_threshold = UVM_PERF_ACCESS_COUNTER_THRESHOLD_DEFAULT;
 
 // Module parameters for the tunables
@@ -100,10 +94,6 @@ MODULE_PARM_DESC(uvm_perf_access_counter_momc_migration_enable,
                  "Whether MOMC access counters will trigger migrations."
                  "Valid values: <= -1 (default policy), 0 (off), >= 1 (on)");
 module_param(uvm_perf_access_counter_batch_count, uint, S_IRUGO);
-module_param(uvm_perf_access_counter_granularity, charp, S_IRUGO);
-MODULE_PARM_DESC(uvm_perf_access_counter_granularity,
-                 "Size of the physical memory region tracked by each counter. Valid values as"
-                 "of Volta: 64k, 2m, 16m, 16g");
 module_param(uvm_perf_access_counter_threshold, uint, S_IRUGO);
 MODULE_PARM_DESC(uvm_perf_access_counter_threshold,
                  "Number of remote accesses on a region required to trigger a notification."
@@ -136,7 +126,7 @@ static va_space_access_counters_info_t *va_space_access_counters_info_get(uvm_va
 
 // Whether access counter migrations are enabled or not. The policy is as
 // follows:
-// - MIMC migrations are enabled by default on P9 systems with ATS support
+// - MIMC migrations are disabled by default on all non-ATS systems.
 // - MOMC migrations are disabled by default on all systems
 // - Users can override this policy by specifying on/off
 static bool is_migration_enabled(uvm_access_counter_type_t type)
@@ -159,7 +149,10 @@ static bool is_migration_enabled(uvm_access_counter_type_t type)
     if (type == UVM_ACCESS_COUNTER_TYPE_MOMC)
         return false;
 
-    return g_uvm_global.ats.supported;
+    if (UVM_ATS_SUPPORTED())
+        return g_uvm_global.ats.supported;
+
+    return false;
 }
 
 // Create the access counters tracking struct for the given VA space
@@ -225,30 +218,18 @@ static NV_STATUS config_granularity_to_bytes(UVM_ACCESS_COUNTER_GRANULARITY gran
     return NV_OK;
 }
 
-// Clear the given access counter and add it to the per-GPU clear tracker
-static NV_STATUS access_counter_clear_targeted(uvm_gpu_t *gpu,
-                                               const uvm_access_counter_buffer_entry_t *entry)
+// Clear the access counter notifications and add it to the per-GPU clear
+// tracker.
+static NV_STATUS access_counter_clear_notifications(uvm_gpu_t *gpu,
+                                                    uvm_access_counter_buffer_entry_t **notification_start,
+                                                    NvU32 num_notifications)
 {
+    NvU32 i;
     NV_STATUS status;
     uvm_push_t push;
     uvm_access_counter_buffer_info_t *access_counters = &gpu->parent->access_counter_buffer_info;
 
-    if (entry->address.is_virtual) {
-        status = uvm_push_begin(gpu->channel_manager,
-                                UVM_CHANNEL_TYPE_MEMOPS,
-                                &push,
-                                "Clear access counter with virtual address: 0x%llx",
-                                entry->address.address);
-    }
-    else {
-        status = uvm_push_begin(gpu->channel_manager,
-                                UVM_CHANNEL_TYPE_MEMOPS,
-                                &push,
-                                "Clear access counter with physical address: 0x%llx:%s",
-                                entry->address.address,
-                                uvm_aperture_string(entry->address.aperture));
-    }
-
+    status = uvm_push_begin(gpu->channel_manager, UVM_CHANNEL_TYPE_MEMOPS, &push, "Clear access counter batch");
     if (status != NV_OK) {
         UVM_ERR_PRINT("Error creating push to clear access counters: %s, GPU %s\n",
                       nvstatusToString(status),
@@ -256,7 +237,8 @@ static NV_STATUS access_counter_clear_targeted(uvm_gpu_t *gpu,
         return status;
     }
 
-    gpu->parent->host_hal->access_counter_clear_targeted(&push, entry);
+    for (i = 0; i < num_notifications; i++)
+        gpu->parent->host_hal->access_counter_clear_targeted(&push, notification_start[i]);
 
     uvm_push_end(&push);
 
@@ -381,25 +363,6 @@ NV_STATUS uvm_gpu_init_access_counters(uvm_parent_gpu_t *parent_gpu)
         g_uvm_access_counter_threshold = uvm_perf_access_counter_threshold;
     }
 
-    if (strcmp(uvm_perf_access_counter_granularity, "64k") == 0) {
-        g_uvm_access_counter_granularity = UVM_ACCESS_COUNTER_GRANULARITY_64K;
-    }
-    else if (strcmp(uvm_perf_access_counter_granularity, "2m") == 0) {
-        g_uvm_access_counter_granularity = UVM_ACCESS_COUNTER_GRANULARITY_2M;
-    }
-    else if (strcmp(uvm_perf_access_counter_granularity, "16m") == 0) {
-        g_uvm_access_counter_granularity = UVM_ACCESS_COUNTER_GRANULARITY_16M;
-    }
-    else if (strcmp(uvm_perf_access_counter_granularity, "16g") == 0) {
-        g_uvm_access_counter_granularity = UVM_ACCESS_COUNTER_GRANULARITY_16G;
-    }
-    else {
-        g_uvm_access_counter_granularity = UVM_ACCESS_COUNTER_GRANULARITY_2M;
-        pr_info("Invalid value '%s' for uvm_perf_access_counter_granularity, using '%s' instead",
-                uvm_perf_access_counter_granularity,
-                UVM_PERF_ACCESS_COUNTER_GRANULARITY_DEFAULT);
-    }
-
     uvm_assert_mutex_locked(&g_uvm_global.global_lock);
     UVM_ASSERT(parent_gpu->access_counter_buffer_hal != NULL);
 
@@ -422,7 +385,7 @@ NV_STATUS uvm_gpu_init_access_counters(uvm_parent_gpu_t *parent_gpu)
     UVM_ASSERT(access_counters->rm_info.bufferSize %
                parent_gpu->access_counter_buffer_hal->entry_size(parent_gpu) == 0);
 
-    status = config_granularity_to_bytes(g_uvm_access_counter_granularity, &granularity_bytes);
+    status = config_granularity_to_bytes(UVM_PERF_ACCESS_COUNTER_GRANULARITY, &granularity_bytes);
     UVM_ASSERT(status == NV_OK);
     if (granularity_bytes > UVM_MAX_TRANSLATION_SIZE)
         UVM_ASSERT(granularity_bytes % UVM_MAX_TRANSLATION_SIZE == 0);
@@ -641,8 +604,8 @@ NV_STATUS uvm_gpu_access_counters_enable(uvm_gpu_t *gpu, uvm_va_space_t *va_spac
     else {
         UvmGpuAccessCntrConfig default_config =
         {
-            .mimcGranularity = g_uvm_access_counter_granularity,
-            .momcGranularity = g_uvm_access_counter_granularity,
+            .mimcGranularity = UVM_PERF_ACCESS_COUNTER_GRANULARITY,
+            .momcGranularity = UVM_PERF_ACCESS_COUNTER_GRANULARITY,
             .mimcUseLimit = UVM_ACCESS_COUNTER_USE_LIMIT_FULL,
             .momcUseLimit = UVM_ACCESS_COUNTER_USE_LIMIT_FULL,
             .threshold = g_uvm_access_counter_threshold,
@@ -717,7 +680,10 @@ static void access_counter_buffer_flush_locked(uvm_gpu_t *gpu, uvm_gpu_buffer_fl
 
     while (get != put) {
         // Wait until valid bit is set
-        UVM_SPIN_WHILE(!gpu->parent->access_counter_buffer_hal->entry_is_valid(gpu->parent, get), &spin);
+        UVM_SPIN_WHILE(!gpu->parent->access_counter_buffer_hal->entry_is_valid(gpu->parent, get), &spin) {
+            if (uvm_global_get_status() != NV_OK)
+                goto done;
+        }
 
         gpu->parent->access_counter_buffer_hal->entry_clear_valid(gpu->parent, get);
         ++get;
@@ -725,6 +691,7 @@ static void access_counter_buffer_flush_locked(uvm_gpu_t *gpu, uvm_gpu_buffer_fl
             get = 0;
     }
 
+done:
     write_get(gpu->parent, get);
 }
 
@@ -767,6 +734,22 @@ static int cmp_sort_virt_notifications_by_instance_ptr(const void *_a, const voi
     return cmp_access_counter_instance_ptr(a, b);
 }
 
+// Sort comparator for pointers to GVA access counter notification buffer
+// entries that sorts by va_space, and fault address.
+static int cmp_sort_virt_notifications_by_va_space_address(const void *_a, const void *_b)
+{
+    const uvm_access_counter_buffer_entry_t **a = (const uvm_access_counter_buffer_entry_t **)_a;
+    const uvm_access_counter_buffer_entry_t **b = (const uvm_access_counter_buffer_entry_t **)_b;
+
+    int result;
+
+    result = UVM_CMP_DEFAULT((*a)->virtual_info.va_space, (*b)->virtual_info.va_space);
+    if (result != 0)
+        return result;
+
+    return UVM_CMP_DEFAULT((*a)->address.address, (*b)->address.address);
+}
+
 // Sort comparator for pointers to GPA access counter notification buffer
 // entries that sorts by physical address' aperture
 static int cmp_sort_phys_notifications_by_processor_id(const void *_a, const void *_b)
@@ -834,12 +817,18 @@ static NvU32 fetch_access_counter_buffer_entries(uvm_gpu_t *gpu,
            (fetch_mode == NOTIFICATION_FETCH_MODE_ALL || notification_index < access_counters->max_batch_size)) {
         uvm_access_counter_buffer_entry_t *current_entry = &notification_cache[notification_index];
 
-        // We cannot just wait for the last entry (the one pointed by put) to become valid, we have to do it
-        // individually since entries can be written out of order
+        // We cannot just wait for the last entry (the one pointed by put) to
+        // become valid, we have to do it individually since entries can be
+        // written out of order
         UVM_SPIN_WHILE(!gpu->parent->access_counter_buffer_hal->entry_is_valid(gpu->parent, get), &spin) {
             // We have some entry to work on. Let's do the rest later.
             if (fetch_mode != NOTIFICATION_FETCH_MODE_ALL && notification_index > 0)
                 goto done;
+
+            // There's no entry to work on and something has gone wrong. Ignore
+            // the rest.
+            if (uvm_global_get_status() != NV_OK)
+               goto done;
         }
 
         // Prevent later accesses being moved above the read of the valid bit
@@ -924,12 +913,11 @@ static void translate_virt_notifications_instance_ptrs(uvm_gpu_t *gpu,
 
 // GVA notifications provide an instance_ptr and ve_id that can be directly
 // translated to a VA space. In order to minimize translations, we sort the
-// entries by instance_ptr.
+// entries by instance_ptr, va_space and notification address in that order.
 static void preprocess_virt_notifications(uvm_gpu_t *gpu,
                                           uvm_access_counter_service_batch_context_t *batch_context)
 {
     if (!batch_context->virt.is_single_instance_ptr) {
-        // Sort by instance_ptr
         sort(batch_context->virt.notifications,
              batch_context->virt.num_notifications,
              sizeof(*batch_context->virt.notifications),
@@ -938,6 +926,12 @@ static void preprocess_virt_notifications(uvm_gpu_t *gpu,
     }
 
     translate_virt_notifications_instance_ptrs(gpu, batch_context);
+
+    sort(batch_context->virt.notifications,
+         batch_context->virt.num_notifications,
+         sizeof(*batch_context->virt.notifications),
+         cmp_sort_virt_notifications_by_va_space_address,
+         NULL);
 }
 
 // GPA notifications provide a physical address and an aperture. Sort
@@ -946,7 +940,6 @@ static void preprocess_virt_notifications(uvm_gpu_t *gpu,
 static void preprocess_phys_notifications(uvm_access_counter_service_batch_context_t *batch_context)
 {
     if (!batch_context->phys.is_single_aperture) {
-        // Sort by instance_ptr
         sort(batch_context->phys.notifications,
              batch_context->phys.num_notifications,
              sizeof(*batch_context->phys.notifications),
@@ -955,6 +948,28 @@ static void preprocess_phys_notifications(uvm_access_counter_service_batch_conte
     }
 }
 
+static NV_STATUS notify_tools_and_process_flags(uvm_gpu_t *gpu,
+                                                uvm_access_counter_buffer_entry_t **notification_start,
+                                                NvU32 num_entries,
+                                                NvU32 flags)
+{
+    NV_STATUS status = NV_OK;
+
+    if (uvm_enable_builtin_tests) {
+        // TODO: Bug 4310744: [UVM][TOOLS] Attribute access counter tools events
+        //                    to va_space instead of broadcasting.
+        NvU32 i;
+
+        for (i = 0; i < num_entries; i++)
+            uvm_tools_broadcast_access_counter(gpu, notification_start[i], flags & UVM_ACCESS_COUNTER_PHYS_ON_MANAGED);
+    }
+
+    if (flags & UVM_ACCESS_COUNTER_ACTION_CLEAR)
+        status = access_counter_clear_notifications(gpu, notification_start, num_entries);
+
+    return status;
+}
+
 static NV_STATUS service_va_block_locked(uvm_processor_id_t processor,
                                          uvm_va_block_t *va_block,
                                          uvm_va_block_retry_t *va_block_retry,
@@ -1009,6 +1024,7 @@ static NV_STATUS service_va_block_locked(uvm_processor_id_t processor,
         NvU64 address = uvm_va_block_cpu_page_address(va_block, page_index);
         bool read_duplicate = false;
         uvm_processor_id_t new_residency;
+        const uvm_va_policy_t *policy;
 
         // Ensure that the migratability iterator covers the current address
         while (iter.end < address)
@@ -1035,21 +1051,23 @@ static NV_STATUS service_va_block_locked(uvm_processor_id_t processor,
 
         // If the underlying VMA is gone, skip HMM migrations.
         if (uvm_va_block_is_hmm(va_block)) {
-            status = uvm_hmm_find_vma(&service_context->block_context, address);
+            status = uvm_hmm_find_vma(service_context->block_context.mm,
+                                      &service_context->block_context.hmm.vma,
+                                      address);
             if (status == NV_ERR_INVALID_ADDRESS)
                 continue;
 
             UVM_ASSERT(status == NV_OK);
         }
 
-        service_context->block_context.policy = uvm_va_policy_get(va_block, address);
+        policy = uvm_va_policy_get(va_block, address);
 
         new_residency = uvm_va_block_select_residency(va_block,
                                                       &service_context->block_context,
                                                       page_index,
                                                       processor,
                                                       uvm_fault_access_type_mask_bit(UVM_FAULT_ACCESS_TYPE_PREFETCH),
-                                                      service_context->block_context.policy,
+                                                      policy,
                                                       &thrashing_hint,
                                                       UVM_SERVICE_OPERATION_ACCESS_COUNTERS,
                                                       &read_duplicate);
@@ -1094,12 +1112,17 @@ static NV_STATUS service_va_block_locked(uvm_processor_id_t processor,
         if (!uvm_processor_mask_empty(&service_context->resident_processors)) {
             while (first_page_index <= last_page_index) {
                 uvm_page_index_t outer = last_page_index + 1;
+                const uvm_va_policy_t *policy;
 
                 if (uvm_va_block_is_hmm(va_block)) {
-                    status = uvm_hmm_find_policy_vma_and_outer(va_block,
-                                                               &service_context->block_context,
-                                                               first_page_index,
-                                                               &outer);
+                    status = NV_ERR_INVALID_ADDRESS;
+                    if (service_context->block_context.mm) {
+                        status = uvm_hmm_find_policy_vma_and_outer(va_block,
+                                                                   &service_context->block_context.hmm.vma,
+                                                                   first_page_index,
+                                                                   &policy,
+                                                                   &outer);
+                    }
                     if (status != NV_OK)
                         break;
                 }
@@ -1155,7 +1178,7 @@ static NV_STATUS service_phys_single_va_block(uvm_gpu_t *gpu,
                                               const uvm_access_counter_buffer_entry_t *current_entry,
                                               const uvm_reverse_map_t *reverse_mappings,
                                               size_t num_reverse_mappings,
-                                              unsigned *out_flags)
+                                              NvU32 *out_flags)
 {
     size_t index;
     uvm_va_block_t *va_block = reverse_mappings[0].va_block;
@@ -1182,7 +1205,6 @@ static NV_STATUS service_phys_single_va_block(uvm_gpu_t *gpu,
         // If an mm is registered with the VA space, we have to retain it
         // in order to lock it before locking the VA space.
         mm = uvm_va_space_mm_retain_lock(va_space);
-
         uvm_va_space_down_read(va_space);
 
         // Re-check that the VA block is valid after taking the VA block lock.
@@ -1243,7 +1265,7 @@ static NV_STATUS service_phys_va_blocks(uvm_gpu_t *gpu,
                                         const uvm_access_counter_buffer_entry_t *current_entry,
                                         const uvm_reverse_map_t *reverse_mappings,
                                         size_t num_reverse_mappings,
-                                        unsigned *out_flags)
+                                        NvU32 *out_flags)
 {
     NV_STATUS status = NV_OK;
     size_t index;
@@ -1251,7 +1273,7 @@ static NV_STATUS service_phys_va_blocks(uvm_gpu_t *gpu,
     *out_flags &= ~UVM_ACCESS_COUNTER_ACTION_CLEAR;
 
     for (index = 0; index < num_reverse_mappings; ++index) {
-        unsigned out_flags_local = 0;
+        NvU32 out_flags_local = 0;
         status = service_phys_single_va_block(gpu,
                                               batch_context,
                                               current_entry,
@@ -1310,7 +1332,7 @@ static NV_STATUS service_phys_notification_translation(uvm_gpu_t *gpu,
                                                        NvU64 address,
                                                        unsigned long sub_granularity,
                                                        size_t *num_reverse_mappings,
-                                                       unsigned *out_flags)
+                                                       NvU32 *out_flags)
 {
     NV_STATUS status;
     NvU32 region_start, region_end;
@@ -1319,7 +1341,10 @@ static NV_STATUS service_phys_notification_translation(uvm_gpu_t *gpu,
 
     // Get the reverse_map translations for all the regions set in the
     // sub_granularity field of the counter.
-    for_each_sub_granularity_region(region_start, region_end, sub_granularity, config->sub_granularity_regions_per_translation) {
+    for_each_sub_granularity_region(region_start,
+                                    region_end,
+                                    sub_granularity,
+                                    config->sub_granularity_regions_per_translation) {
         NvU64 local_address = address + region_start * config->sub_granularity_region_size;
         NvU32 local_translation_size = (region_end - region_start) * config->sub_granularity_region_size;
         uvm_reverse_map_t *local_reverse_mappings = batch_context->phys.translations + *num_reverse_mappings;
@@ -1368,7 +1393,7 @@ static NV_STATUS service_phys_notification_translation(uvm_gpu_t *gpu,
 static NV_STATUS service_phys_notification(uvm_gpu_t *gpu,
                                            uvm_access_counter_service_batch_context_t *batch_context,
                                            const uvm_access_counter_buffer_entry_t *current_entry,
-                                           unsigned *out_flags)
+                                           NvU32 *out_flags)
 {
     NvU64 address;
     NvU64 translation_index;
@@ -1379,7 +1404,7 @@ static NV_STATUS service_phys_notification(uvm_gpu_t *gpu,
     size_t total_reverse_mappings = 0;
     uvm_gpu_t *resident_gpu = NULL;
     NV_STATUS status = NV_OK;
-    unsigned flags = 0;
+    NvU32 flags = 0;
 
     address = current_entry->address.address;
     UVM_ASSERT(address % config->translation_size == 0);
@@ -1407,7 +1432,7 @@ static NV_STATUS service_phys_notification(uvm_gpu_t *gpu,
 
     for (translation_index = 0; translation_index < config->translations_per_counter; ++translation_index) {
         size_t num_reverse_mappings;
-        unsigned out_flags_local = 0;
+        NvU32 out_flags_local = 0;
         status = service_phys_notification_translation(gpu,
                                                        resident_gpu,
                                                        batch_context,
@@ -1429,11 +1454,8 @@ static NV_STATUS service_phys_notification(uvm_gpu_t *gpu,
         sub_granularity = sub_granularity >> config->sub_granularity_regions_per_translation;
     }
 
-    // Currently we only report events for our tests, not for tools
-    if (uvm_enable_builtin_tests) {
-        *out_flags |= UVM_ACCESS_COUNTER_ACTION_NOTIFY;
-        *out_flags |= ((total_reverse_mappings != 0) ? UVM_ACCESS_COUNTER_ON_MANAGED : 0);
-    }
+    if (uvm_enable_builtin_tests)
+        *out_flags |= ((total_reverse_mappings != 0) ? UVM_ACCESS_COUNTER_PHYS_ON_MANAGED : 0);
 
     if (status == NV_OK && (flags & UVM_ACCESS_COUNTER_ACTION_CLEAR))
         *out_flags |= UVM_ACCESS_COUNTER_ACTION_CLEAR;
@@ -1446,22 +1468,21 @@ static NV_STATUS service_phys_notifications(uvm_gpu_t *gpu,
                                             uvm_access_counter_service_batch_context_t *batch_context)
 {
     NvU32 i;
+    uvm_access_counter_buffer_entry_t **notifications = batch_context->phys.notifications;
+
     preprocess_phys_notifications(batch_context);
 
     for (i = 0; i < batch_context->phys.num_notifications; ++i) {
         NV_STATUS status;
-        uvm_access_counter_buffer_entry_t *current_entry = batch_context->phys.notifications[i];
-        unsigned flags = 0;
+        uvm_access_counter_buffer_entry_t *current_entry = notifications[i];
+        NvU32 flags = 0;
 
         if (!UVM_ID_IS_VALID(current_entry->physical_info.resident_id))
             continue;
 
         status = service_phys_notification(gpu, batch_context, current_entry, &flags);
-        if (flags & UVM_ACCESS_COUNTER_ACTION_NOTIFY)
-            uvm_tools_broadcast_access_counter(gpu, current_entry, flags & UVM_ACCESS_COUNTER_ON_MANAGED);
 
-        if (status == NV_OK && (flags & UVM_ACCESS_COUNTER_ACTION_CLEAR))
-            status = access_counter_clear_targeted(gpu, current_entry);
+        notify_tools_and_process_flags(gpu, &notifications[i], 1, flags);
 
         if (status != NV_OK)
             return status;
@@ -1470,187 +1491,375 @@ static NV_STATUS service_phys_notifications(uvm_gpu_t *gpu,
     return NV_OK;
 }
 
-static int cmp_sort_gpu_phys_addr(const void *_a, const void *_b)
+static NV_STATUS service_notification_va_block_helper(struct mm_struct *mm,
+                                                      uvm_va_block_t *va_block,
+                                                      uvm_processor_id_t processor,
+                                                      uvm_access_counter_service_batch_context_t *batch_context)
 {
-    return uvm_gpu_phys_addr_cmp(*(uvm_gpu_phys_address_t*)_a,
-                                 *(uvm_gpu_phys_address_t*)_b);
-}
+    uvm_va_block_retry_t va_block_retry;
+    uvm_page_mask_t *accessed_pages = &batch_context->accessed_pages;
+    uvm_service_block_context_t *service_context = &batch_context->block_service_context;
 
-static bool gpu_phys_same_region(uvm_gpu_phys_address_t a, uvm_gpu_phys_address_t b, NvU64 granularity)
-{
-    if (a.aperture != b.aperture)
-        return false;
-
-    UVM_ASSERT(is_power_of_2(granularity));
-
-    return UVM_ALIGN_DOWN(a.address, granularity) == UVM_ALIGN_DOWN(b.address, granularity);
-}
-
-static bool phys_address_in_accessed_sub_region(uvm_gpu_phys_address_t address,
-                                                NvU64 region_size,
-                                                NvU64 sub_region_size,
-                                                NvU32 accessed_mask)
-{
-    const unsigned accessed_index = (address.address % region_size) / sub_region_size;
-
-    // accessed_mask is only filled for tracking granularities larger than 64K
-    if (region_size == UVM_PAGE_SIZE_64K)
-        return true;
-
-    UVM_ASSERT(accessed_index < 32);
-    return ((1 << accessed_index) & accessed_mask) != 0;
-}
-
-static NV_STATUS service_virt_notification(uvm_gpu_t *gpu,
-                                           uvm_access_counter_service_batch_context_t *batch_context,
-                                           const uvm_access_counter_buffer_entry_t *current_entry,
-                                           unsigned *out_flags)
-{
-    NV_STATUS status = NV_OK;
-    NvU64 notification_size;
-    NvU64 address;
-    uvm_processor_id_t *resident_processors = batch_context->virt.scratch.resident_processors;
-    uvm_gpu_phys_address_t *phys_addresses = batch_context->virt.scratch.phys_addresses;
-    int num_addresses = 0;
-    int i;
-
-    // Virtual address notifications are always 64K aligned
-    NvU64 region_start = current_entry->address.address;
-    NvU64 region_end = current_entry->address.address + UVM_PAGE_SIZE_64K;
-    
-
-    uvm_access_counter_buffer_info_t *access_counters = &gpu->parent->access_counter_buffer_info;
-    uvm_access_counter_type_t counter_type = current_entry->counter_type;
-
-    const uvm_gpu_access_counter_type_config_t *config = get_config_for_type(access_counters, counter_type);
-
-    uvm_va_space_t *va_space = current_entry->virtual_info.va_space;
-
-    UVM_ASSERT(counter_type == UVM_ACCESS_COUNTER_TYPE_MIMC);
-
-    // Entries with NULL va_space are simply dropped.
-    if (!va_space)
+    if (uvm_page_mask_empty(accessed_pages))
         return NV_OK;
 
-    status = config_granularity_to_bytes(config->rm.granularity, &notification_size);
-    if (status != NV_OK)
-        return status;
+    uvm_assert_mutex_locked(&va_block->lock);
 
-    // Collect physical locations that could have been touched
-    // in the reported 64K VA region. The notification mask can
-    // correspond to any of them.
-    uvm_va_space_down_read(va_space);
-    for (address = region_start; address < region_end;) {
-        uvm_va_block_t *va_block;
+    service_context->operation = UVM_SERVICE_OPERATION_ACCESS_COUNTERS;
+    service_context->num_retries = 0;
+    service_context->block_context.mm = mm;
 
-        NV_STATUS local_status = uvm_va_block_find(va_space, address, &va_block);
-        if (local_status == NV_ERR_INVALID_ADDRESS || local_status == NV_ERR_OBJECT_NOT_FOUND) {
-            address += PAGE_SIZE;
-            continue;
-        }
+    return UVM_VA_BLOCK_RETRY_LOCKED(va_block,
+                                     &va_block_retry,
+                                     service_va_block_locked(processor,
+                                                             va_block,
+                                                             &va_block_retry,
+                                                             service_context,
+                                                             accessed_pages));
+}
 
-        uvm_mutex_lock(&va_block->lock);
-        while (address < va_block->end && address < region_end) {
-            const unsigned page_index = uvm_va_block_cpu_page_index(va_block, address);
+static void expand_notification_block(uvm_gpu_va_space_t *gpu_va_space,
+                                      uvm_va_block_t *va_block,
+                                      uvm_page_mask_t *accessed_pages,
+                                      const uvm_access_counter_buffer_entry_t *current_entry)
+{
+    NvU64 addr;
+    NvU64 granularity = 0;
+    uvm_gpu_t *resident_gpu = NULL;
+    uvm_processor_id_t resident_id;
+    uvm_page_index_t page_index;
+    uvm_gpu_t *gpu = gpu_va_space->gpu;
+    const uvm_access_counter_buffer_info_t *access_counters = &gpu->parent->access_counter_buffer_info;
+    const uvm_gpu_access_counter_type_config_t *config = get_config_for_type(access_counters,
+                                                                             UVM_ACCESS_COUNTER_TYPE_MIMC);
 
-            // UVM va_block always maps the closest resident location to processor
-            const uvm_processor_id_t res_id = uvm_va_block_page_get_closest_resident(va_block, page_index, gpu->id);
+    config_granularity_to_bytes(config->rm.granularity, &granularity);
 
-            // Add physical location if it's valid and not local vidmem
-            if (UVM_ID_IS_VALID(res_id) && !uvm_id_equal(res_id, gpu->id)) {
-                uvm_gpu_phys_address_t phys_address = uvm_va_block_res_phys_page_address(va_block, page_index, res_id, gpu);
-                if (phys_address_in_accessed_sub_region(phys_address,
-                                                        notification_size,
-                                                        config->sub_granularity_region_size,
-                                                        current_entry->sub_granularity)) {
-                    resident_processors[num_addresses] = res_id;
-                    phys_addresses[num_addresses] = phys_address;
-                    ++num_addresses;
-                }
-                else {
-                    UVM_DBG_PRINT_RL("Skipping phys address %llx:%s, because it couldn't have been accessed in mask %x",
-                                     phys_address.address,
-                                     uvm_aperture_string(phys_address.aperture),
-                                     current_entry->sub_granularity);
-                }
-            }
+    // Granularities other than 2MB can only be enabled by UVM tests. Do nothing
+    // in that case.
+    if (granularity != UVM_PAGE_SIZE_2M)
+        return;
 
-            address += PAGE_SIZE;
-        }
-        uvm_mutex_unlock(&va_block->lock);
+    addr = current_entry->address.address;
+
+    uvm_assert_rwsem_locked(&gpu_va_space->va_space->lock);
+    uvm_assert_mutex_locked(&va_block->lock);
+
+    page_index = uvm_va_block_cpu_page_index(va_block, addr);
+
+    resident_id = uvm_va_block_page_get_closest_resident(va_block, page_index, gpu->id);
+
+    // resident_id might be invalid or might already be the same as the GPU
+    // which received the notification if the memory was already migrated before
+    // acquiring the locks either during the servicing of previous notifications
+    // or during faults or because of explicit migrations or if the VA range was
+    // freed after receiving the notification. Return NV_OK in such cases.
+    if (!UVM_ID_IS_VALID(resident_id) || uvm_id_equal(resident_id, gpu->id))
+        return;
+
+    if (UVM_ID_IS_GPU(resident_id))
+        resident_gpu = uvm_va_space_get_gpu(gpu_va_space->va_space, resident_id);
+
+    if (uvm_va_block_get_physical_size(va_block, resident_id, page_index) != granularity) {
+        uvm_page_mask_set(accessed_pages, page_index);
     }
-    uvm_va_space_up_read(va_space);
+    else {
+        NvU32 region_start;
+        NvU32 region_end;
+        unsigned long sub_granularity = current_entry->sub_granularity;
+        NvU32 num_regions = config->sub_granularity_regions_per_translation;
+        NvU32 num_sub_pages = config->sub_granularity_region_size / PAGE_SIZE;
+        uvm_page_mask_t *resident_mask = uvm_va_block_resident_mask_get(va_block, resident_id);
 
-    // The addresses need to be sorted to aid coalescing.
-    sort(phys_addresses,
-         num_addresses,
-         sizeof(*phys_addresses),
-         cmp_sort_gpu_phys_addr,
-         NULL);
+        UVM_ASSERT(num_sub_pages >= 1);
 
-    for (i = 0; i < num_addresses; ++i) {
-        uvm_access_counter_buffer_entry_t *fake_entry = &batch_context->virt.scratch.phys_entry;
-
-        // Skip the current pointer if the physical region was already handled
-        if (i > 0 && gpu_phys_same_region(phys_addresses[i - 1], phys_addresses[i], notification_size)) {
-            UVM_ASSERT(uvm_id_equal(resident_processors[i - 1], resident_processors[i]));
-            continue;
+        // region_start and region_end refer to sub_granularity indices, not
+        // page_indices.
+        for_each_sub_granularity_region(region_start, region_end, sub_granularity, num_regions) {
+            uvm_page_mask_region_fill(accessed_pages,
+                                      uvm_va_block_region(region_start * num_sub_pages,
+                                                          region_end * num_sub_pages));
         }
-        UVM_DBG_PRINT_RL("Faking MIMC address[%i/%i]: %llx (granularity mask: %llx) in aperture %s on device %s\n",
-                         i,
-                         num_addresses,
-                         phys_addresses[i].address,
-                         notification_size - 1,
-                         uvm_aperture_string(phys_addresses[i].aperture),
-                         uvm_gpu_name(gpu));
 
-        // Construct a fake phys addr AC entry
-        fake_entry->counter_type = current_entry->counter_type;
-        fake_entry->address.address = UVM_ALIGN_DOWN(phys_addresses[i].address, notification_size);
-        fake_entry->address.aperture = phys_addresses[i].aperture;
-        fake_entry->address.is_virtual = false;
-        fake_entry->physical_info.resident_id = resident_processors[i];
-        fake_entry->counter_value = current_entry->counter_value;
-        fake_entry->sub_granularity = current_entry->sub_granularity;
+        // Remove pages in the va_block which are not resident on resident_id.
+        // If the GPU is heavily accessing those pages, future access counter
+        // migrations will migrate them to the GPU.
+        uvm_page_mask_and(accessed_pages, accessed_pages, resident_mask);
+    }
+}
 
-        status = service_phys_notification(gpu, batch_context, fake_entry, out_flags);
-        if (status != NV_OK)
+static NV_STATUS service_virt_notifications_in_block(uvm_gpu_va_space_t *gpu_va_space,
+                                                     struct mm_struct *mm,
+                                                     uvm_va_block_t *va_block,
+                                                     uvm_access_counter_service_batch_context_t *batch_context,
+                                                     NvU32 index,
+                                                     NvU32 *out_index)
+{
+    NvU32 i;
+    NvU32 flags = 0;
+    NV_STATUS status = NV_OK;
+    NV_STATUS flags_status;
+    uvm_gpu_t *gpu = gpu_va_space->gpu;
+    uvm_va_space_t *va_space = gpu_va_space->va_space;
+    uvm_page_mask_t *accessed_pages = &batch_context->accessed_pages;
+    uvm_access_counter_buffer_entry_t **notifications = batch_context->virt.notifications;
+
+    UVM_ASSERT(va_block);
+    UVM_ASSERT(index < batch_context->virt.num_notifications);
+
+    uvm_assert_rwsem_locked(&va_space->lock);
+
+    uvm_page_mask_zero(accessed_pages);
+
+    uvm_mutex_lock(&va_block->lock);
+
+    for (i = index; i < batch_context->virt.num_notifications; i++) {
+        uvm_access_counter_buffer_entry_t *current_entry = notifications[i];
+        NvU64 address = current_entry->address.address;
+
+        if ((current_entry->virtual_info.va_space == va_space) && (address <= va_block->end))
+            expand_notification_block(gpu_va_space, va_block, accessed_pages, current_entry);
+        else
             break;
     }
 
+    *out_index = i;
+
+    // Atleast one notification should have been processed.
+    UVM_ASSERT(index < *out_index);
+
+    status = service_notification_va_block_helper(mm, va_block, gpu->id, batch_context);
+
+    uvm_mutex_unlock(&va_block->lock);
+
+    if (status == NV_OK)
+        flags |= UVM_ACCESS_COUNTER_ACTION_CLEAR;
+
+    flags_status = notify_tools_and_process_flags(gpu, &notifications[index], *out_index - index, flags);
+
+    if ((status == NV_OK) && (flags_status != NV_OK))
+        status = flags_status;
+
+    return status;
+}
+
+static NV_STATUS service_virt_notification_ats(uvm_gpu_va_space_t *gpu_va_space,
+                                               struct mm_struct *mm,
+                                               uvm_access_counter_service_batch_context_t *batch_context,
+                                               NvU32 index,
+                                               NvU32 *out_index)
+{
+
+    NvU32 i;
+    NvU64 base;
+    NvU64 end;
+    NvU64 address;
+    NvU32 flags = UVM_ACCESS_COUNTER_ACTION_CLEAR;
+    NV_STATUS status = NV_OK;
+    NV_STATUS flags_status;
+    struct vm_area_struct *vma = NULL;
+    uvm_gpu_t *gpu = gpu_va_space->gpu;
+    uvm_va_space_t *va_space = gpu_va_space->va_space;
+    uvm_ats_fault_context_t *ats_context = &batch_context->ats_context;
+    uvm_access_counter_buffer_entry_t **notifications = batch_context->virt.notifications;
+
+    UVM_ASSERT(index < batch_context->virt.num_notifications);
+
+    uvm_assert_mmap_lock_locked(mm);
+    uvm_assert_rwsem_locked(&va_space->lock);
+
+    address = notifications[index]->address.address;
+
+    vma = find_vma_intersection(mm, address, address + 1);
+    if (!vma) {
+        // Clear the notification entry to continue receiving access counter
+        // notifications when a new VMA is allocated in this range.
+        status = notify_tools_and_process_flags(gpu, &notifications[index], 1, flags);
+        *out_index = index + 1;
+        return status;
+    }
+
+    base = UVM_VA_BLOCK_ALIGN_DOWN(address);
+    end = min(base + UVM_VA_BLOCK_SIZE, (NvU64)vma->vm_end);
+
+    uvm_page_mask_zero(&ats_context->accessed_mask);
+
+    for (i = index; i < batch_context->virt.num_notifications; i++) {
+        uvm_access_counter_buffer_entry_t *current_entry = notifications[i];
+        address = current_entry->address.address;
+
+        if ((current_entry->virtual_info.va_space == va_space) && (address < end))
+            uvm_page_mask_set(&ats_context->accessed_mask, (address - base) / PAGE_SIZE);
+        else
+            break;
+    }
+
+    *out_index = i;
+
+    // Atleast one notification should have been processed.
+    UVM_ASSERT(index < *out_index);
+
+    // TODO: Bug 2113632: [UVM] Don't clear access counters when the preferred
+    //                    location is set
+    // If no pages were actually migrated, don't clear the access counters.
+    status = uvm_ats_service_access_counters(gpu_va_space, vma, base, ats_context);
+    if (status != NV_OK)
+        flags &= ~UVM_ACCESS_COUNTER_ACTION_CLEAR;
+
+    flags_status = notify_tools_and_process_flags(gpu, &notifications[index], *out_index - index, flags);
+    if ((status == NV_OK) && (flags_status != NV_OK))
+        status = flags_status;
+
+    return status;
+}
+
+static NV_STATUS service_virt_notifications_batch(uvm_gpu_va_space_t *gpu_va_space,
+                                                  struct mm_struct *mm,
+                                                  uvm_access_counter_service_batch_context_t *batch_context,
+                                                  NvU32 index,
+                                                  NvU32 *out_index)
+{
+    NV_STATUS status;
+    uvm_va_range_t *va_range;
+    uvm_va_space_t *va_space = gpu_va_space->va_space;
+    uvm_access_counter_buffer_entry_t *current_entry = batch_context->virt.notifications[index];
+    NvU64 address = current_entry->address.address;
+
+    UVM_ASSERT(va_space);
+
+    if (mm)
+        uvm_assert_mmap_lock_locked(mm);
+
+    uvm_assert_rwsem_locked(&va_space->lock);
+
+    // Virtual address notifications are always 64K aligned
+    UVM_ASSERT(IS_ALIGNED(address, UVM_PAGE_SIZE_64K));
+
+    va_range = uvm_va_range_find(va_space, address);
+    if (va_range) {
+        // Avoid clearing the entry by default.
+        NvU32 flags = 0;
+        uvm_va_block_t *va_block = NULL;
+
+        if (va_range->type == UVM_VA_RANGE_TYPE_MANAGED) {
+            size_t index = uvm_va_range_block_index(va_range, address);
+
+            va_block = uvm_va_range_block(va_range, index);
+
+            // If the va_range is a managed range, the notification belongs to a
+            // recently freed va_range if va_block is NULL. If va_block is not
+            // NULL, service_virt_notifications_in_block will process flags.
+            // Clear the notification entry to continue receiving notifications
+            // when a new va_range is allocated in that region.
+            flags = UVM_ACCESS_COUNTER_ACTION_CLEAR;
+        }
+
+        if (va_block) {
+            status = service_virt_notifications_in_block(gpu_va_space, mm, va_block, batch_context, index, out_index);
+        }
+        else {
+            status = notify_tools_and_process_flags(gpu_va_space->gpu, batch_context->virt.notifications, 1, flags);
+            *out_index = index + 1;
+        }
+    }
+    else if (uvm_ats_can_service_faults(gpu_va_space, mm)) {
+        status = service_virt_notification_ats(gpu_va_space, mm, batch_context, index, out_index);
+    }
+    else {
+        NvU32 flags;
+        uvm_va_block_t *va_block = NULL;
+
+        status = uvm_hmm_va_block_find(va_space, address, &va_block);
+
+        // TODO: Bug 4309292: [UVM][HMM] Re-enable access counter HMM block
+        //                    migrations for virtual notifications
+        //
+        // - If the va_block is HMM, don't clear the notification since HMM
+        // migrations are currently disabled.
+        //
+        // - If the va_block isn't HMM, the notification belongs to a recently
+        // freed va_range. Clear the notification entry to continue receiving
+        // notifications when a new va_range is allocated in this region.
+        flags = va_block ? 0 : UVM_ACCESS_COUNTER_ACTION_CLEAR;
+
+        UVM_ASSERT((status == NV_ERR_OBJECT_NOT_FOUND) ||
+                   (status == NV_ERR_INVALID_ADDRESS)  ||
+                   uvm_va_block_is_hmm(va_block));
+
+        // Clobber status to continue processing the rest of the notifications
+        // in the batch.
+        status = notify_tools_and_process_flags(gpu_va_space->gpu, batch_context->virt.notifications, 1, flags);
+
+        *out_index = index + 1;
+    }
+
     return status;
 }
 
 static NV_STATUS service_virt_notifications(uvm_gpu_t *gpu,
                                             uvm_access_counter_service_batch_context_t *batch_context)
 {
-    NvU32 i;
+    NvU32 i = 0;
     NV_STATUS status = NV_OK;
+    struct mm_struct *mm = NULL;
+    uvm_va_space_t *va_space = NULL;
+    uvm_va_space_t *prev_va_space = NULL;
+    uvm_gpu_va_space_t *gpu_va_space = NULL;
+
+    // TODO: Bug 4299018 : Add support for virtual access counter migrations on
+    //                     4K page sizes.
+    if (PAGE_SIZE == UVM_PAGE_SIZE_4K) {
+        return notify_tools_and_process_flags(gpu,
+                                              batch_context->virt.notifications,
+                                              batch_context->virt.num_notifications,
+                                              0);
+    }
+
     preprocess_virt_notifications(gpu, batch_context);
 
-    for (i = 0; i < batch_context->virt.num_notifications; ++i) {
-        unsigned flags = 0;
+    while (i < batch_context->virt.num_notifications) {
         uvm_access_counter_buffer_entry_t *current_entry = batch_context->virt.notifications[i];
+        va_space = current_entry->virtual_info.va_space;
 
-        status = service_virt_notification(gpu, batch_context, current_entry, &flags);
+        if (va_space != prev_va_space) {
 
-        UVM_DBG_PRINT_RL("Processed virt access counter (%d/%d): %sMANAGED (status: %d) clear: %s\n",
-                         i + 1,
-                         batch_context->virt.num_notifications,
-                         (flags & UVM_ACCESS_COUNTER_ON_MANAGED) ? "" : "NOT ",
-                         status,
-                         (flags & UVM_ACCESS_COUNTER_ACTION_CLEAR) ? "YES" : "NO");
+            // New va_space detected, drop locks of the old va_space.
+            if (prev_va_space) {
+                uvm_va_space_up_read(prev_va_space);
+                uvm_va_space_mm_release_unlock(prev_va_space, mm);
 
-        if (uvm_enable_builtin_tests)
-            uvm_tools_broadcast_access_counter(gpu, current_entry, flags & UVM_ACCESS_COUNTER_ON_MANAGED);
+                mm = NULL;
+                gpu_va_space = NULL;
+            }
 
-        if (status == NV_OK && (flags & UVM_ACCESS_COUNTER_ACTION_CLEAR))
-            status = access_counter_clear_targeted(gpu, current_entry);
+            // Acquire locks for the new va_space.
+            if (va_space) {
+                mm = uvm_va_space_mm_retain_lock(va_space);
+                uvm_va_space_down_read(va_space);
+
+                gpu_va_space = uvm_gpu_va_space_get_by_parent_gpu(va_space, gpu->parent);
+            }
+
+            prev_va_space = va_space;
+        }
+
+        if (va_space && gpu_va_space && uvm_va_space_has_access_counter_migrations(va_space)) {
+            status = service_virt_notifications_batch(gpu_va_space, mm, batch_context, i, &i);
+        }
+        else {
+            status = notify_tools_and_process_flags(gpu, &batch_context->virt.notifications[i], 1, 0);
+            i++;
+        }
 
         if (status != NV_OK)
             break;
     }
 
+    if (va_space) {
+        uvm_va_space_up_read(va_space);
+        uvm_va_space_mm_release_unlock(va_space, mm);
+    }
+
     return status;
 }
 
@@ -1933,6 +2142,7 @@ NV_STATUS uvm_test_reset_access_counters(UVM_TEST_RESET_ACCESS_COUNTERS_PARAMS *
     }
     else {
         uvm_access_counter_buffer_entry_t entry = { 0 };
+        uvm_access_counter_buffer_entry_t *notification = &entry;
 
         if (params->counter_type == UVM_TEST_ACCESS_COUNTER_TYPE_MIMC)
             entry.counter_type = UVM_ACCESS_COUNTER_TYPE_MIMC;
@@ -1942,7 +2152,7 @@ NV_STATUS uvm_test_reset_access_counters(UVM_TEST_RESET_ACCESS_COUNTERS_PARAMS *
         entry.bank = params->bank;
         entry.tag = params->tag;
 
-        status = access_counter_clear_targeted(gpu, &entry);
+        status = access_counter_clear_notifications(gpu, &notification, 1);
     }
 
     if (status == NV_OK)

kernel-open/nvidia-uvm/uvm_gpu_non_replayable_faults.c

@@ -235,17 +235,27 @@ static NV_STATUS fetch_non_replayable_fault_buffer_entries(uvm_parent_gpu_t *par
     return NV_OK;
 }
 
-// In SRIOV, the UVM (guest) driver does not have access to the privileged
-// registers used to clear the faulted bit. Instead, UVM requests host RM to do
-// the clearing on its behalf, using a SW method.
 static bool use_clear_faulted_channel_sw_method(uvm_gpu_t *gpu)
 {
-    if (uvm_gpu_is_virt_mode_sriov(gpu)) {
-        UVM_ASSERT(gpu->parent->has_clear_faulted_channel_sw_method);
-        return true;
-    }
+    // If true, UVM uses a SW method to request RM to do the clearing on its
+    // behalf.
+    bool use_sw_method = false;
 
-    return false;
+    // In SRIOV, the UVM (guest) driver does not have access to the privileged
+    // registers used to clear the faulted bit.
+    if (uvm_gpu_is_virt_mode_sriov(gpu))
+        use_sw_method = true;
+
+    // In Confidential Computing access to the privileged registers is blocked,
+    // in order to prevent interference between guests, or between the
+    // (untrusted) host and the guests.
+    if (g_uvm_global.conf_computing_enabled)
+        use_sw_method = true;
+
+    if (use_sw_method)
+        UVM_ASSERT(gpu->parent->has_clear_faulted_channel_sw_method);
+
+    return use_sw_method;
 }
 
 static NV_STATUS clear_faulted_method_on_gpu(uvm_gpu_t *gpu,
@@ -343,6 +353,7 @@ static NV_STATUS service_managed_fault_in_block_locked(uvm_gpu_t *gpu,
     bool read_duplicate;
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
     uvm_non_replayable_fault_buffer_info_t *non_replayable_faults = &gpu->parent->fault_buffer_info.non_replayable;
+    const uvm_va_policy_t *policy;
 
     UVM_ASSERT(!fault_entry->is_fatal);
 
@@ -352,7 +363,7 @@ static NV_STATUS service_managed_fault_in_block_locked(uvm_gpu_t *gpu,
     UVM_ASSERT(fault_entry->fault_address >= va_block->start);
     UVM_ASSERT(fault_entry->fault_address <= va_block->end);
 
-    service_context->block_context.policy = uvm_va_policy_get(va_block, fault_entry->fault_address);
+    policy = uvm_va_policy_get(va_block, fault_entry->fault_address);
 
     if (service_context->num_retries == 0) {
         // notify event to tools/performance heuristics. For now we use a
@@ -361,7 +372,7 @@ static NV_STATUS service_managed_fault_in_block_locked(uvm_gpu_t *gpu,
         uvm_perf_event_notify_gpu_fault(&va_space->perf_events,
                                         va_block,
                                         gpu->id,
-                                        service_context->block_context.policy->preferred_location,
+                                        policy->preferred_location,
                                         fault_entry,
                                         ++non_replayable_faults->batch_id,
                                         false);
@@ -396,7 +407,7 @@ static NV_STATUS service_managed_fault_in_block_locked(uvm_gpu_t *gpu,
                                                   page_index,
                                                   gpu->id,
                                                   fault_entry->access_type_mask,
-                                                  service_context->block_context.policy,
+                                                  policy,
                                                   &thrashing_hint,
                                                   UVM_SERVICE_OPERATION_NON_REPLAYABLE_FAULTS,
                                                   &read_duplicate);
@@ -569,7 +580,7 @@ static NV_STATUS service_non_managed_fault(uvm_gpu_va_space_t *gpu_va_space,
 
         ats_context->client_type = UVM_FAULT_CLIENT_TYPE_HUB;
 
-        ats_invalidate->write_faults_in_batch = false;
+        ats_invalidate->tlb_batch_pending = false;
 
         va_range_next = uvm_va_space_iter_first(gpu_va_space->va_space, fault_entry->fault_address, ~0ULL);
 
@@ -678,10 +689,17 @@ static NV_STATUS service_fault(uvm_gpu_t *gpu, uvm_fault_buffer_entry_t *fault_e
     fault_entry->fault_source.channel_id = user_channel->hw_channel_id;
 
     if (!fault_entry->is_fatal) {
-        status = uvm_va_block_find_create(fault_entry->va_space,
-                                          fault_entry->fault_address,
-                                          va_block_context,
-                                          &va_block);
+        if (mm) {
+            status = uvm_va_block_find_create(fault_entry->va_space,
+                                              fault_entry->fault_address,
+                                              &va_block_context->hmm.vma,
+                                              &va_block);
+        }
+        else {
+            status = uvm_va_block_find_create_managed(fault_entry->va_space,
+                                                      fault_entry->fault_address,
+                                                      &va_block);
+        }
         if (status == NV_OK)
             status = service_managed_fault_in_block(gpu_va_space->gpu, va_block, fault_entry);
         else
@@ -734,8 +752,6 @@ void uvm_gpu_service_non_replayable_fault_buffer(uvm_gpu_t *gpu)
         // Differently to replayable faults, we do not batch up and preprocess
         // non-replayable faults since getting multiple faults on the same
         // memory region is not very likely
-        //
-        // TODO: Bug 2103669: [UVM/ATS] Optimize ATS fault servicing
         for (i = 0; i < cached_faults; ++i) {
             status = service_fault(gpu, &gpu->parent->fault_buffer_info.non_replayable.fault_cache[i]);
             if (status != NV_OK)

kernel-open/nvidia-uvm/uvm_gpu_replayable_faults.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2022 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -362,7 +362,8 @@ static NV_STATUS push_cancel_on_gpu(uvm_gpu_t *gpu,
                                         "Cancel targeting instance_ptr {0x%llx:%s}\n",
                                         instance_ptr.address,
                                         uvm_aperture_string(instance_ptr.aperture));
-    } else {
+    }
+    else {
         status = uvm_push_begin_acquire(gpu->channel_manager,
                                         UVM_CHANNEL_TYPE_MEMOPS,
                                         &replayable_faults->replay_tracker,
@@ -623,7 +624,15 @@ static NV_STATUS fault_buffer_flush_locked(uvm_gpu_t *gpu,
 
     while (get != put) {
         // Wait until valid bit is set
-        UVM_SPIN_WHILE(!parent_gpu->fault_buffer_hal->entry_is_valid(parent_gpu, get), &spin);
+        UVM_SPIN_WHILE(!parent_gpu->fault_buffer_hal->entry_is_valid(parent_gpu, get), &spin) {
+            // Channels might be idle (e.g. in teardown) so check for errors
+            // actively. In that case the gpu pointer is valid.
+            NV_STATUS status = gpu ? uvm_channel_manager_check_errors(gpu->channel_manager) : uvm_global_get_status();
+            if (status != NV_OK) {
+                write_get(parent_gpu, get);
+                return status;
+            }
+        }
 
         fault_buffer_skip_replayable_entry(parent_gpu, get);
         ++get;
@@ -697,9 +706,6 @@ static inline int cmp_access_type(uvm_fault_access_type_t a, uvm_fault_access_ty
 
 typedef enum
 {
-    // Fetch a batch of faults from the buffer.
-    FAULT_FETCH_MODE_BATCH_ALL,
-
     // Fetch a batch of faults from the buffer. Stop at the first entry that is
     // not ready yet
     FAULT_FETCH_MODE_BATCH_READY,
@@ -857,9 +863,11 @@ static NV_STATUS fetch_fault_buffer_entries(uvm_gpu_t *gpu,
         // written out of order
         UVM_SPIN_WHILE(!gpu->parent->fault_buffer_hal->entry_is_valid(gpu->parent, get), &spin) {
             // We have some entry to work on. Let's do the rest later.
-            if (fetch_mode != FAULT_FETCH_MODE_ALL &&
-                fetch_mode != FAULT_FETCH_MODE_BATCH_ALL &&
-                fault_index > 0)
+            if (fetch_mode == FAULT_FETCH_MODE_BATCH_READY && fault_index > 0)
+                goto done;
+            
+            status = uvm_global_get_status();
+            if (status != NV_OK)
                 goto done;
         }
 
@@ -888,6 +896,7 @@ static NV_STATUS fetch_fault_buffer_entries(uvm_gpu_t *gpu,
 
         current_entry->va_space = NULL;
         current_entry->filtered = false;
+        current_entry->replayable.cancel_va_mode = UVM_FAULT_CANCEL_VA_MODE_ALL;
 
         if (current_entry->fault_source.utlb_id > batch_context->max_utlb_id) {
             UVM_ASSERT(current_entry->fault_source.utlb_id < replayable_faults->utlb_count);
@@ -1184,7 +1193,11 @@ static void mark_fault_fatal(uvm_fault_service_batch_context_t *batch_context,
     fault_entry->replayable.cancel_va_mode = cancel_va_mode;
 
     utlb->has_fatal_faults = true;
-    batch_context->has_fatal_faults = true;
+
+    if (!batch_context->fatal_va_space) {
+        UVM_ASSERT(fault_entry->va_space);
+        batch_context->fatal_va_space = fault_entry->va_space;
+    }
 }
 
 static void fault_entry_duplicate_flags(uvm_fault_service_batch_context_t *batch_context,
@@ -1322,6 +1335,7 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
     uvm_fault_buffer_entry_t **ordered_fault_cache = batch_context->ordered_fault_cache;
     uvm_service_block_context_t *block_context = &replayable_faults->block_service_context;
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
+    const uvm_va_policy_t *policy;
     NvU64 end;
 
     // Check that all uvm_fault_access_type_t values can fit into an NvU8
@@ -1347,13 +1361,13 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
     UVM_ASSERT(ordered_fault_cache[first_fault_index]->fault_address <= va_block->end);
 
     if (uvm_va_block_is_hmm(va_block)) {
-        uvm_hmm_find_policy_end(va_block,
-                                &block_context->block_context,
-                                ordered_fault_cache[first_fault_index]->fault_address,
-                                &end);
+        policy = uvm_hmm_find_policy_end(va_block,
+                                         block_context->block_context.hmm.vma,
+                                         ordered_fault_cache[first_fault_index]->fault_address,
+                                         &end);
     }
     else {
-        block_context->block_context.policy = uvm_va_range_get_policy(va_block->va_range);
+        policy = uvm_va_range_get_policy(va_block->va_range);
         end = va_block->end;
     }
 
@@ -1377,7 +1391,10 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
         UVM_ASSERT(current_entry->fault_access_type ==
                    uvm_fault_access_type_mask_highest(current_entry->access_type_mask));
 
-        current_entry->is_fatal            = false;
+        // Unserviceable faults were already skipped by the caller. There are no
+        // unserviceable fault types that could be in the same VA block as a
+        // serviceable fault.
+        UVM_ASSERT(!current_entry->is_fatal);
         current_entry->is_throttled        = false;
         current_entry->is_invalid_prefetch = false;
 
@@ -1393,7 +1410,7 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
             update_batch_and_notify_fault(gpu,
                                           batch_context,
                                           va_block,
-                                          block_context->block_context.policy->preferred_location,
+                                          policy->preferred_location,
                                           current_entry,
                                           is_duplicate);
         }
@@ -1473,7 +1490,7 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
                                                       page_index,
                                                       gpu->id,
                                                       service_access_type_mask,
-                                                      block_context->block_context.policy,
+                                                      policy,
                                                       &thrashing_hint,
                                                       UVM_SERVICE_OPERATION_REPLAYABLE_FAULTS,
                                                       &read_duplicate);
@@ -1511,7 +1528,7 @@ static NV_STATUS service_fault_batch_block_locked(uvm_gpu_t *gpu,
 
     ++block_context->num_retries;
 
-    if (status == NV_OK && batch_context->has_fatal_faults)
+    if (status == NV_OK && batch_context->fatal_va_space)
         status = uvm_va_block_set_cancel(va_block, &block_context->block_context, gpu);
 
     return status;
@@ -1625,21 +1642,25 @@ static NV_STATUS service_fault_batch_ats_sub_vma(uvm_gpu_va_space_t *gpu_va_spac
     uvm_ats_fault_context_t *ats_context = &batch_context->ats_context;
     const uvm_page_mask_t *read_fault_mask = &ats_context->read_fault_mask;
     const uvm_page_mask_t *write_fault_mask = &ats_context->write_fault_mask;
-    const uvm_page_mask_t *faults_serviced_mask = &ats_context->faults_serviced_mask;
     const uvm_page_mask_t *reads_serviced_mask = &ats_context->reads_serviced_mask;
-    uvm_page_mask_t *tmp_mask = &ats_context->tmp_mask;
+    uvm_page_mask_t *faults_serviced_mask = &ats_context->faults_serviced_mask;
+    uvm_page_mask_t *accessed_mask = &ats_context->accessed_mask;
 
     UVM_ASSERT(vma);
 
     ats_context->client_type = UVM_FAULT_CLIENT_TYPE_GPC;
 
-    uvm_page_mask_or(tmp_mask, write_fault_mask, read_fault_mask);
+    uvm_page_mask_or(accessed_mask, write_fault_mask, read_fault_mask);
 
     status = uvm_ats_service_faults(gpu_va_space, vma, base, &batch_context->ats_context);
 
-    UVM_ASSERT(uvm_page_mask_subset(faults_serviced_mask, tmp_mask));
+    // Remove prefetched pages from the serviced mask since fault servicing
+    // failures belonging to prefetch pages need to be ignored.
+    uvm_page_mask_and(faults_serviced_mask, faults_serviced_mask, accessed_mask);
 
-    if ((status != NV_OK) || uvm_page_mask_equal(faults_serviced_mask, tmp_mask)) {
+    UVM_ASSERT(uvm_page_mask_subset(faults_serviced_mask, accessed_mask));
+
+    if ((status != NV_OK) || uvm_page_mask_equal(faults_serviced_mask, accessed_mask)) {
         (*block_faults) += (fault_index_end - fault_index_start);
         return status;
     }
@@ -1671,7 +1692,8 @@ static NV_STATUS service_fault_batch_ats_sub_vma(uvm_gpu_va_space_t *gpu_va_spac
         if (access_type <= UVM_FAULT_ACCESS_TYPE_READ) {
             cancel_va_mode = UVM_FAULT_CANCEL_VA_MODE_ALL;
         }
-        else if (access_type >= UVM_FAULT_ACCESS_TYPE_WRITE) {
+	else {
+            UVM_ASSERT(access_type >= UVM_FAULT_ACCESS_TYPE_WRITE);
             if (uvm_fault_access_type_mask_test(current_entry->access_type_mask, UVM_FAULT_ACCESS_TYPE_READ) &&
                 !uvm_page_mask_test(reads_serviced_mask, page_index))
                 cancel_va_mode = UVM_FAULT_CANCEL_VA_MODE_ALL;
@@ -1730,6 +1752,10 @@ static NV_STATUS service_fault_batch_ats_sub(uvm_gpu_va_space_t *gpu_va_space,
         uvm_fault_access_type_t access_type = current_entry->fault_access_type;
         bool is_duplicate = check_fault_entry_duplicate(current_entry, previous_entry);
 
+        // ATS faults can't be unserviceable, since unserviceable faults require
+        // GMMU PTEs.
+        UVM_ASSERT(!current_entry->is_fatal);
+
         i++;
 
         update_batch_and_notify_fault(gpu_va_space->gpu,
@@ -1867,7 +1893,13 @@ static NV_STATUS service_fault_batch_dispatch(uvm_va_space_t *va_space,
         va_range_next = uvm_va_space_iter_next(va_range_next, ~0ULL);
     }
 
-    status = uvm_va_block_find_create_in_range(va_space, va_range, fault_address, va_block_context, &va_block);
+    if (va_range)
+        status = uvm_va_block_find_create_in_range(va_space, va_range, fault_address, &va_block);
+    else if (mm)
+        status = uvm_hmm_va_block_find_create(va_space, fault_address, &va_block_context->hmm.vma, &va_block);
+    else
+        status = NV_ERR_INVALID_ADDRESS;
+
     if (status == NV_OK) {
         status = service_fault_batch_block(gpu, va_block, batch_context, fault_index, block_faults);
     }
@@ -1923,14 +1955,198 @@ static NV_STATUS service_fault_batch_dispatch(uvm_va_space_t *va_space,
     return status;
 }
 
+// Called when a fault in the batch has been marked fatal. Flush the buffer
+// under the VA and mmap locks to remove any potential stale fatal faults, then
+// service all new faults for just that VA space and cancel those which are
+// fatal. Faults in other VA spaces are replayed when done and will be processed
+// when normal fault servicing resumes.
+static NV_STATUS service_fault_batch_for_cancel(uvm_gpu_t *gpu, uvm_fault_service_batch_context_t *batch_context)
+{
+    NV_STATUS status = NV_OK;
+    NvU32 i;
+    uvm_va_space_t *va_space = batch_context->fatal_va_space;
+    uvm_gpu_va_space_t *gpu_va_space = NULL;
+    struct mm_struct *mm;
+    uvm_replayable_fault_buffer_info_t *replayable_faults = &gpu->parent->fault_buffer_info.replayable;
+    uvm_service_block_context_t *service_context = &gpu->parent->fault_buffer_info.replayable.block_service_context;
+    uvm_va_block_context_t *va_block_context = &service_context->block_context;
+
+    UVM_ASSERT(gpu->parent->replayable_faults_supported);
+    UVM_ASSERT(va_space);
+
+    // Perform the flush and re-fetch while holding the mmap_lock and the
+    // VA space lock. This avoids stale faults because it prevents any vma
+    // modifications (mmap, munmap, mprotect) from happening between the time HW
+    // takes the fault and we cancel it.
+    mm = uvm_va_space_mm_retain_lock(va_space);
+    va_block_context->mm = mm;
+    uvm_va_space_down_read(va_space);
+
+    // We saw fatal faults in this VA space before. Flush while holding
+    // mmap_lock to make sure those faults come back (aren't stale).
+    //
+    // We need to wait until all old fault messages have arrived before
+    // flushing, hence UVM_GPU_BUFFER_FLUSH_MODE_WAIT_UPDATE_PUT.
+    status = fault_buffer_flush_locked(gpu,
+                                       UVM_GPU_BUFFER_FLUSH_MODE_WAIT_UPDATE_PUT,
+                                       UVM_FAULT_REPLAY_TYPE_START,
+                                       batch_context);
+    if (status != NV_OK)
+        goto done;
+
+    // Wait for the flush's replay to finish to give the legitimate faults a
+    // chance to show up in the buffer again.
+    status = uvm_tracker_wait(&replayable_faults->replay_tracker);
+    if (status != NV_OK)
+        goto done;
+
+    // We expect all replayed faults to have arrived in the buffer so we can re-
+    // service them. The replay-and-wait sequence above will ensure they're all
+    // in the HW buffer. When GSP owns the HW buffer, we also have to wait for
+    // GSP to copy all available faults from the HW buffer into the shadow
+    // buffer.
+    //
+    // TODO: Bug 2533557: This flush does not actually guarantee that GSP will
+    //       copy over all faults.
+    status = hw_fault_buffer_flush_locked(gpu->parent);
+    if (status != NV_OK)
+        goto done;
+
+    // If there is no GPU VA space for the GPU, ignore all faults in the VA
+    // space. This can happen if the GPU VA space has been destroyed since we
+    // unlocked the VA space in service_fault_batch. That means the fatal faults
+    // are stale, because unregistering the GPU VA space requires preempting the
+    // context and detaching all channels in that VA space. Restart fault
+    // servicing from the top.
+    gpu_va_space = uvm_gpu_va_space_get_by_parent_gpu(va_space, gpu->parent);
+    if (!gpu_va_space)
+        goto done;
+
+    // Re-parse the new faults
+    batch_context->num_invalid_prefetch_faults = 0;
+    batch_context->num_duplicate_faults        = 0;
+    batch_context->num_replays                 = 0;
+    batch_context->fatal_va_space              = NULL;
+    batch_context->has_throttled_faults        = false;
+
+    status = fetch_fault_buffer_entries(gpu, batch_context, FAULT_FETCH_MODE_ALL);
+    if (status != NV_OK)
+        goto done;
+
+    // No more faults left. Either the previously-seen fatal entry was stale, or
+    // RM killed the context underneath us.
+    if (batch_context->num_cached_faults == 0)
+        goto done;
+
+    ++batch_context->batch_id;
+
+    status = preprocess_fault_batch(gpu, batch_context);
+    if (status != NV_OK) {
+        if (status == NV_WARN_MORE_PROCESSING_REQUIRED) {
+            // Another flush happened due to stale faults or a context-fatal
+            // error. The previously-seen fatal fault might not exist anymore,
+            // so restart fault servicing from the top.
+            status = NV_OK;
+        }
+
+        goto done;
+    }
+
+    // Search for the target VA space
+    for (i = 0; i < batch_context->num_coalesced_faults; i++) {
+        uvm_fault_buffer_entry_t *current_entry = batch_context->ordered_fault_cache[i];
+        UVM_ASSERT(current_entry->va_space);
+        if (current_entry->va_space == va_space)
+            break;
+    }
+
+    while (i < batch_context->num_coalesced_faults) {
+        uvm_fault_buffer_entry_t *current_entry = batch_context->ordered_fault_cache[i];
+
+        if (current_entry->va_space != va_space)
+            break;
+
+        // service_fault_batch_dispatch() doesn't expect unserviceable faults.
+        // Just cancel them directly.
+        if (current_entry->is_fatal) {
+            status = cancel_fault_precise_va(gpu, current_entry, UVM_FAULT_CANCEL_VA_MODE_ALL);
+            if (status != NV_OK)
+                break;
+
+            ++i;
+        }
+        else {
+            uvm_ats_fault_invalidate_t *ats_invalidate = &gpu->parent->fault_buffer_info.replayable.ats_invalidate;
+            NvU32 block_faults;
+
+            ats_invalidate->tlb_batch_pending = false;
+            uvm_hmm_service_context_init(service_context);
+
+            // Service all the faults that we can. We only really need to search
+            // for fatal faults, but attempting to service all is the easiest
+            // way to do that.
+            status = service_fault_batch_dispatch(va_space, gpu_va_space, batch_context, i, &block_faults, false);
+            if (status != NV_OK) {
+                // TODO: Bug 3900733: clean up locking in service_fault_batch().
+                // We need to drop lock and retry. That means flushing and
+                // starting over.
+                if (status == NV_WARN_MORE_PROCESSING_REQUIRED)
+                    status = NV_OK;
+
+                break;
+            }
+
+            // Invalidate TLBs before cancel to ensure that fatal faults don't
+            // get stuck in HW behind non-fatal faults to the same line.
+            status = uvm_ats_invalidate_tlbs(gpu_va_space, ats_invalidate, &batch_context->tracker);
+            if (status != NV_OK)
+                break;
+
+            while (block_faults-- > 0) {
+                current_entry = batch_context->ordered_fault_cache[i];
+                if (current_entry->is_fatal) {
+                    status = cancel_fault_precise_va(gpu, current_entry, current_entry->replayable.cancel_va_mode);
+                    if (status != NV_OK)
+                        break;
+                }
+
+                ++i;
+            }
+        }
+    }
+
+done:
+    uvm_va_space_up_read(va_space);
+    uvm_va_space_mm_release_unlock(va_space, mm);
+
+    if (status == NV_OK) {
+        // There are two reasons to flush the fault buffer here.
+        //
+        // 1) Functional. We need to replay both the serviced non-fatal faults
+        //    and the skipped faults in other VA spaces. The former need to be
+        //    restarted and the latter need to be replayed so the normal fault
+        //    service mechanism can fetch and process them.
+        //
+        // 2) Performance. After cancelling the fatal faults, a flush removes
+        //    any potential duplicated fault that may have been added while
+        //    processing the faults in this batch. This flush also avoids doing
+        //    unnecessary processing after the fatal faults have been cancelled,
+        //    so all the rest are unlikely to remain after a replay because the
+        //    context is probably in the process of dying.
+        status = fault_buffer_flush_locked(gpu,
+                                           UVM_GPU_BUFFER_FLUSH_MODE_UPDATE_PUT,
+                                           UVM_FAULT_REPLAY_TYPE_START,
+                                           batch_context);
+    }
+
+    return status;
+}
 // Scan the ordered view of faults and group them by different va_blocks
 // (managed faults) and service faults for each va_block, in batch.
 // Service non-managed faults one at a time as they are encountered during the
 // scan.
 //
-// This function returns NV_WARN_MORE_PROCESSING_REQUIRED if the fault buffer
-// was flushed because the needs_fault_buffer_flush flag was set on some GPU VA
-// space
+// Fatal faults are marked for later processing by the caller.
 static NV_STATUS service_fault_batch(uvm_gpu_t *gpu,
                                      fault_service_mode_t service_mode,
                                      uvm_fault_service_batch_context_t *batch_context)
@@ -1949,7 +2165,7 @@ static NV_STATUS service_fault_batch(uvm_gpu_t *gpu,
 
     UVM_ASSERT(gpu->parent->replayable_faults_supported);
 
-    ats_invalidate->write_faults_in_batch = false;
+    ats_invalidate->tlb_batch_pending = false;
     uvm_hmm_service_context_init(service_context);
 
     for (i = 0; i < batch_context->num_coalesced_faults;) {
@@ -1984,38 +2200,25 @@ static NV_STATUS service_fault_batch(uvm_gpu_t *gpu,
             va_block_context->mm = mm;
 
             uvm_va_space_down_read(va_space);
-
             gpu_va_space = uvm_gpu_va_space_get_by_parent_gpu(va_space, gpu->parent);
-            if (uvm_processor_mask_test_and_clear_atomic(&va_space->needs_fault_buffer_flush, gpu->id)) {
-                status = fault_buffer_flush_locked(gpu,
-                                                   UVM_GPU_BUFFER_FLUSH_MODE_WAIT_UPDATE_PUT,
-                                                   UVM_FAULT_REPLAY_TYPE_START,
-                                                   batch_context);
-                if (status == NV_OK)
-                    status = NV_WARN_MORE_PROCESSING_REQUIRED;
-
-                break;
-            }
-
-            // The case where there is no valid GPU VA space for the GPU in this
-            // VA space is handled next
         }
 
         // Some faults could be already fatal if they cannot be handled by
         // the UVM driver
         if (current_entry->is_fatal) {
             ++i;
-            batch_context->has_fatal_faults = true;
+            if (!batch_context->fatal_va_space)
+                batch_context->fatal_va_space = va_space;
+
             utlb->has_fatal_faults = true;
             UVM_ASSERT(utlb->num_pending_faults > 0);
             continue;
         }
 
-        if (!uvm_processor_mask_test(&va_space->registered_gpu_va_spaces, gpu->parent->id)) {
+        if (!gpu_va_space) {
             // If there is no GPU VA space for the GPU, ignore the fault. This
             // can happen if a GPU VA space is destroyed without explicitly
-            // freeing all memory ranges (destroying the VA range triggers a
-            // flush of the fault buffer) and there are stale entries in the
+            // freeing all memory ranges and there are stale entries in the
             // buffer that got fixed by the servicing in a previous batch.
             ++i;
             continue;
@@ -2033,15 +2236,17 @@ static NV_STATUS service_fault_batch(uvm_gpu_t *gpu,
             uvm_va_space_mm_release_unlock(va_space, mm);
             mm = NULL;
             va_space = NULL;
+            status = NV_OK;
             continue;
         }
+
         if (status != NV_OK)
             goto fail;
 
         i += block_faults;
 
         // Don't issue replays in cancel mode
-        if (replay_per_va_block && !batch_context->has_fatal_faults) {
+        if (replay_per_va_block && !batch_context->fatal_va_space) {
             status = push_replay_on_gpu(gpu, UVM_FAULT_REPLAY_TYPE_START, batch_context);
             if (status != NV_OK)
                 goto fail;
@@ -2053,8 +2258,6 @@ static NV_STATUS service_fault_batch(uvm_gpu_t *gpu,
         }
     }
 
-    // Only clobber status if invalidate_status != NV_OK, since status may also
-    // contain NV_WARN_MORE_PROCESSING_REQUIRED.
     if (va_space != NULL) {
         NV_STATUS invalidate_status = uvm_ats_invalidate_tlbs(gpu_va_space, ats_invalidate, &batch_context->tracker);
         if (invalidate_status != NV_OK)
@@ -2262,77 +2465,48 @@ static NvU32 is_fatal_fault_in_buffer(uvm_fault_service_batch_context_t *batch_c
     return false;
 }
 
-typedef enum
-{
-    // Only cancel faults flagged as fatal
-    FAULT_CANCEL_MODE_FATAL,
-
-    // Cancel all faults in the batch unconditionally
-    FAULT_CANCEL_MODE_ALL,
-} fault_cancel_mode_t;
-
-// Cancel faults in the given fault service batch context. The function provides
-// two different modes depending on the value of cancel_mode:
-// - If cancel_mode == FAULT_CANCEL_MODE_FATAL, only faults flagged as fatal
-// will be cancelled. In this case, the reason reported to tools is the one
-// contained in the fault entry itself.
-// - If cancel_mode == FAULT_CANCEL_MODE_ALL, all faults will be cancelled
-// unconditionally. In this case, the reason reported to tools for non-fatal
-// faults is the one passed to this function.
-static NV_STATUS cancel_faults_precise_va(uvm_gpu_t *gpu,
-                                          uvm_fault_service_batch_context_t *batch_context,
-                                          fault_cancel_mode_t cancel_mode,
-                                          UvmEventFatalReason reason)
+// Cancel all faults in the given fault service batch context, even those not
+// marked as fatal.
+static NV_STATUS cancel_faults_all(uvm_gpu_t *gpu,
+                                   uvm_fault_service_batch_context_t *batch_context,
+                                   UvmEventFatalReason reason)
 {
     NV_STATUS status = NV_OK;
     NV_STATUS fault_status;
-    uvm_va_space_t *va_space = NULL;
-    NvU32 i;
+    NvU32 i = 0;
 
     UVM_ASSERT(gpu->parent->fault_cancel_va_supported);
-    if (cancel_mode == FAULT_CANCEL_MODE_ALL)
-        UVM_ASSERT(reason != UvmEventFatalReasonInvalid);
+    UVM_ASSERT(reason != UvmEventFatalReasonInvalid);
 
-    for (i = 0; i < batch_context->num_coalesced_faults; ++i) {
+    while (i < batch_context->num_coalesced_faults && status == NV_OK) {
         uvm_fault_buffer_entry_t *current_entry = batch_context->ordered_fault_cache[i];
+        uvm_va_space_t *va_space = current_entry->va_space;
+        bool skip_va_space;
 
-        UVM_ASSERT(current_entry->va_space);
+        UVM_ASSERT(va_space);
 
-        if (current_entry->va_space != va_space) {
-            // Fault on a different va_space, drop the lock of the old one...
-            if (va_space != NULL)
-                uvm_va_space_up_read(va_space);
+        uvm_va_space_down_read(va_space);
 
-            va_space = current_entry->va_space;
+        // If there is no GPU VA space for the GPU, ignore all faults in
+        // that VA space. This can happen if the GPU VA space has been
+        // destroyed since we unlocked the VA space in service_fault_batch.
+        // Ignoring the fault avoids targetting a PDB that might have been
+        // reused by another process.
+        skip_va_space = !uvm_gpu_va_space_get_by_parent_gpu(va_space, gpu->parent);
 
-            // ... and take the lock of the new one
-            uvm_va_space_down_read(va_space);
+        for (;
+             i < batch_context->num_coalesced_faults && current_entry->va_space == va_space;
+             current_entry = batch_context->ordered_fault_cache[++i]) {
+            uvm_fault_cancel_va_mode_t cancel_va_mode;
 
-            // We don't need to check whether a buffer flush is required
-            // (due to VA range destruction).
-            // - For cancel_mode == FAULT_CANCEL_MODE_FATAL, once a fault is
-            // flagged as fatal we need to cancel it, even if its VA range no
-            // longer exists.
-            // - For cancel_mode == FAULT_CANCEL_MODE_ALL we don't care about
-            // any of this, we just want to trigger RC in RM.
-        }
+            if (skip_va_space)
+                continue;
 
-        if (!uvm_processor_mask_test(&va_space->registered_gpu_va_spaces, gpu->parent->id)) {
-            // If there is no GPU VA space for the GPU, ignore the fault.
-            // This can happen if the GPU VA did not exist in
-            // service_fault_batch(), or it was destroyed since then.
-            // This is to avoid targetting a PDB that might have been reused
-            // by another process.
-            continue;
-        }
-
-        // Cancel the fault
-        if (cancel_mode == FAULT_CANCEL_MODE_ALL || current_entry->is_fatal) {
-            uvm_fault_cancel_va_mode_t cancel_va_mode = current_entry->replayable.cancel_va_mode;
-
-            // If cancelling unconditionally and the fault was not fatal,
-            // set the cancel reason passed to this function
-            if (!current_entry->is_fatal) {
+            if (current_entry->is_fatal) {
+                UVM_ASSERT(current_entry->fatal_reason != UvmEventFatalReasonInvalid);
+                cancel_va_mode = current_entry->replayable.cancel_va_mode;
+            }
+            else {
                 current_entry->fatal_reason = reason;
                 cancel_va_mode = UVM_FAULT_CANCEL_VA_MODE_ALL;
             }
@@ -2341,17 +2515,13 @@ static NV_STATUS cancel_faults_precise_va(uvm_gpu_t *gpu,
             if (status != NV_OK)
                 break;
         }
+
+        uvm_va_space_up_read(va_space);
     }
 
-    if (va_space != NULL)
-        uvm_va_space_up_read(va_space);
-
-    // After cancelling the fatal faults, the fault buffer is flushed to remove
-    // any potential duplicated fault that may have been added while processing
-    // the faults in this batch. This flush also avoids doing unnecessary
-    // processing after the fatal faults have been cancelled, so all the rest
-    // are unlikely to remain after a replay because the context is probably in
-    // the process of dying.
+    // Because each cancel itself triggers a replay, there may be a large number
+    // of new duplicated faults in the buffer after cancelling all the known
+    // ones. Flushing the buffer discards them to avoid unnecessary processing.
     fault_status = fault_buffer_flush_locked(gpu,
                                              UVM_GPU_BUFFER_FLUSH_MODE_UPDATE_PUT,
                                              UVM_FAULT_REPLAY_TYPE_START,
@@ -2399,12 +2569,12 @@ static void cancel_fault_batch(uvm_gpu_t *gpu,
                                uvm_fault_service_batch_context_t *batch_context,
                                UvmEventFatalReason reason)
 {
-    if (gpu->parent->fault_cancel_va_supported) {
-        cancel_faults_precise_va(gpu, batch_context, FAULT_CANCEL_MODE_ALL, reason);
-        return;
-    }
-
-    cancel_fault_batch_tlb(gpu, batch_context, reason);
+    // Return code is ignored since we're on a global error path and wouldn't be
+    // able to recover anyway.
+    if (gpu->parent->fault_cancel_va_supported)
+        cancel_faults_all(gpu, batch_context, reason);
+    else
+        cancel_fault_batch_tlb(gpu, batch_context, reason);
 }
 
 
@@ -2491,7 +2661,7 @@ static NV_STATUS cancel_faults_precise_tlb(uvm_gpu_t *gpu, uvm_fault_service_bat
 
         batch_context->num_invalid_prefetch_faults = 0;
         batch_context->num_replays                 = 0;
-        batch_context->has_fatal_faults            = false;
+        batch_context->fatal_va_space              = NULL;
         batch_context->has_throttled_faults        = false;
 
         // 5) Fetch all faults from buffer
@@ -2538,9 +2708,6 @@ static NV_STATUS cancel_faults_precise_tlb(uvm_gpu_t *gpu, uvm_fault_service_bat
         // 8) Service all non-fatal faults and mark all non-serviceable faults
         // as fatal
         status = service_fault_batch(gpu, FAULT_SERVICE_MODE_CANCEL, batch_context);
-        if (status == NV_WARN_MORE_PROCESSING_REQUIRED)
-            continue;
-
         UVM_ASSERT(batch_context->num_replays == 0);
         if (status == NV_ERR_NO_MEMORY)
             continue;
@@ -2548,7 +2715,7 @@ static NV_STATUS cancel_faults_precise_tlb(uvm_gpu_t *gpu, uvm_fault_service_bat
             break;
 
         // No more fatal faults left, we are done
-        if (!batch_context->has_fatal_faults)
+        if (!batch_context->fatal_va_space)
             break;
 
         // 9) Search for uTLBs that contain fatal faults and meet the
@@ -2570,13 +2737,9 @@ static NV_STATUS cancel_faults_precise_tlb(uvm_gpu_t *gpu, uvm_fault_service_bat
 
 static NV_STATUS cancel_faults_precise(uvm_gpu_t *gpu, uvm_fault_service_batch_context_t *batch_context)
 {
-    UVM_ASSERT(batch_context->has_fatal_faults);
-    if (gpu->parent->fault_cancel_va_supported) {
-        return cancel_faults_precise_va(gpu,
-                                        batch_context,
-                                        FAULT_CANCEL_MODE_FATAL,
-                                        UvmEventFatalReasonInvalid);
-    }
+    UVM_ASSERT(batch_context->fatal_va_space);
+    if (gpu->parent->fault_cancel_va_supported)
+        return service_fault_batch_for_cancel(gpu, batch_context);
 
     return cancel_faults_precise_tlb(gpu, batch_context);
 }
@@ -2632,7 +2795,7 @@ void uvm_gpu_service_replayable_faults(uvm_gpu_t *gpu)
         batch_context->num_invalid_prefetch_faults = 0;
         batch_context->num_duplicate_faults        = 0;
         batch_context->num_replays                 = 0;
-        batch_context->has_fatal_faults            = false;
+        batch_context->fatal_va_space              = NULL;
         batch_context->has_throttled_faults        = false;
 
         status = fetch_fault_buffer_entries(gpu, batch_context, FAULT_FETCH_MODE_BATCH_READY);
@@ -2660,9 +2823,6 @@ void uvm_gpu_service_replayable_faults(uvm_gpu_t *gpu)
         // was flushed
         num_replays += batch_context->num_replays;
 
-        if (status == NV_WARN_MORE_PROCESSING_REQUIRED)
-            continue;
-
         enable_disable_prefetch_faults(gpu->parent, batch_context);
 
         if (status != NV_OK) {
@@ -2676,10 +2836,17 @@ void uvm_gpu_service_replayable_faults(uvm_gpu_t *gpu)
             break;
         }
 
-        if (batch_context->has_fatal_faults) {
+        if (batch_context->fatal_va_space) {
             status = uvm_tracker_wait(&batch_context->tracker);
-            if (status == NV_OK)
+            if (status == NV_OK) {
                 status = cancel_faults_precise(gpu, batch_context);
+                if (status == NV_OK) {
+                    // Cancel handling should've issued at least one replay
+                    UVM_ASSERT(batch_context->num_replays > 0);
+                    ++num_batches;
+                    continue;
+                }
+            }
 
             break;
         }

kernel-open/nvidia-uvm/uvm_gpu_semaphore.c

@@ -579,8 +579,10 @@ static void uvm_gpu_semaphore_encrypted_payload_update(uvm_channel_t *channel, u
     void *auth_tag_cpu_addr = uvm_rm_mem_get_cpu_va(semaphore->conf_computing.auth_tag);
     NvU32 *gpu_notifier_cpu_addr = (NvU32 *)uvm_rm_mem_get_cpu_va(semaphore->conf_computing.notifier);
     NvU32 *payload_cpu_addr = (NvU32 *)uvm_rm_mem_get_cpu_va(semaphore->conf_computing.encrypted_payload);
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
 
-    UVM_ASSERT(uvm_channel_is_secure_ce(channel));
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(gpu));
+    UVM_ASSERT(uvm_channel_is_ce(channel));
 
     last_observed_notifier = semaphore->conf_computing.last_observed_notifier;
     gpu_notifier = UVM_READ_ONCE(*gpu_notifier_cpu_addr);

kernel-open/nvidia-uvm/uvm_gpu_semaphore.h

@@ -91,9 +91,9 @@ struct uvm_gpu_tracking_semaphore_struct
 // Create a semaphore pool for a GPU.
 NV_STATUS uvm_gpu_semaphore_pool_create(uvm_gpu_t *gpu, uvm_gpu_semaphore_pool_t **pool_out);
 
-// When the Confidential Computing feature is enabled, pools associated with
-// secure CE channels are allocated in the CPR of vidmem and as such have
-// all the associated access restrictions. Because of this, they're called
+// When the Confidential Computing feature is enabled, semaphore pools
+// associated with CE channels are allocated in the CPR of vidmem and as such
+// have all the associated access restrictions. Because of this, they're called
 // secure pools and secure semaphores are allocated out of said secure pools.
 NV_STATUS uvm_gpu_semaphore_secure_pool_create(uvm_gpu_t *gpu, uvm_gpu_semaphore_pool_t **pool_out);
 

kernel-open/nvidia-uvm/uvm_hmm.c

@@ -71,6 +71,24 @@ module_param(uvm_disable_hmm, bool, 0444);
 #include "uvm_va_policy.h"
 #include "uvm_tools.h"
 
+// The function nv_PageSwapCache() wraps the check for page swap cache flag in
+// order to support a wide variety of kernel versions.
+// The function PageSwapCache() is removed after 32f51ead3d77 ("mm: remove
+// PageSwapCache") in v6.12-rc1.
+// The function folio_test_swapcache() was added in Linux 5.16 (d389a4a811551
+// "mm: Add folio flag manipulation functions")
+// Systems with HMM patches backported to 5.14 are possible, but those systems
+// do not include folio_test_swapcache()
+// TODO: Bug 4050579: Remove this when migration of swap cached pages is updated
+static __always_inline bool nv_PageSwapCache(struct page *page)
+{
+#if defined(NV_FOLIO_TEST_SWAPCACHE_PRESENT)
+    return folio_test_swapcache(page_folio(page));
+#else
+    return PageSwapCache(page);
+#endif
+}
+
 static NV_STATUS gpu_chunk_add(uvm_va_block_t *va_block,
                                uvm_page_index_t page_index,
                                struct page *page);
@@ -110,7 +128,20 @@ typedef struct
 
 bool uvm_hmm_is_enabled_system_wide(void)
 {
-    return !uvm_disable_hmm && !g_uvm_global.ats.enabled && uvm_va_space_mm_enabled_system();
+    if (uvm_disable_hmm)
+        return false;
+
+    if (g_uvm_global.ats.enabled)
+        return false;
+
+    // Confidential Computing and HMM impose mutually exclusive constraints. In
+    // Confidential Computing the GPU can only access pages resident in vidmem,
+    // but in HMM pages may be required to be resident in sysmem: file backed
+    // VMAs, huge pages, etc.
+    if (g_uvm_global.conf_computing_enabled)
+        return false;
+
+    return uvm_va_space_mm_enabled_system();
 }
 
 bool uvm_hmm_is_enabled(uvm_va_space_t *va_space)
@@ -127,32 +158,17 @@ static uvm_va_block_t *hmm_va_block_from_node(uvm_range_tree_node_t *node)
     return container_of(node, uvm_va_block_t, hmm.node);
 }
 
-NV_STATUS uvm_hmm_va_space_initialize(uvm_va_space_t *va_space)
+void uvm_hmm_va_space_initialize(uvm_va_space_t *va_space)
 {
     uvm_hmm_va_space_t *hmm_va_space = &va_space->hmm;
-    struct mm_struct *mm = va_space->va_space_mm.mm;
-    int ret;
 
     if (!uvm_hmm_is_enabled(va_space))
-        return NV_OK;
-
-    uvm_assert_mmap_lock_locked_write(mm);
-    uvm_assert_rwsem_locked_write(&va_space->lock);
+        return;
 
     uvm_range_tree_init(&hmm_va_space->blocks);
     uvm_mutex_init(&hmm_va_space->blocks_lock, UVM_LOCK_ORDER_LEAF);
 
-    // Initialize MMU interval notifiers for this process.
-    // This allows mmu_interval_notifier_insert() to be called without holding
-    // the mmap_lock for write.
-    // Note: there is no __mmu_notifier_unregister(), this call just allocates
-    // memory which is attached to the mm_struct and freed when the mm_struct
-    // is freed.
-    ret = __mmu_notifier_register(NULL, mm);
-    if (ret)
-        return errno_to_nv_status(ret);
-
-    return NV_OK;
+    return;
 }
 
 void uvm_hmm_va_space_destroy(uvm_va_space_t *va_space)
@@ -325,7 +341,6 @@ static bool hmm_invalidate(uvm_va_block_t *va_block,
     region = uvm_va_block_region_from_start_end(va_block, start, end);
 
     va_block_context->hmm.vma = NULL;
-    va_block_context->policy = NULL;
 
     // We only need to unmap GPUs since Linux handles the CPUs.
     for_each_gpu_id_in_mask(id, &va_block->mapped) {
@@ -444,11 +459,11 @@ static void hmm_va_block_init(uvm_va_block_t *va_block,
 static NV_STATUS hmm_va_block_find_create(uvm_va_space_t *va_space,
                                           NvU64 addr,
                                           bool allow_unreadable_vma,
-                                          uvm_va_block_context_t *va_block_context,
+                                          struct vm_area_struct **vma_out,
                                           uvm_va_block_t **va_block_ptr)
 {
-    struct mm_struct *mm = va_space->va_space_mm.mm;
-    struct vm_area_struct *vma;
+    struct mm_struct *mm;
+    struct vm_area_struct *va_block_vma;
     uvm_va_block_t *va_block;
     NvU64 start, end;
     NV_STATUS status;
@@ -457,15 +472,14 @@ static NV_STATUS hmm_va_block_find_create(uvm_va_space_t *va_space,
     if (!uvm_hmm_is_enabled(va_space))
         return NV_ERR_INVALID_ADDRESS;
 
-    UVM_ASSERT(mm);
-    UVM_ASSERT(!va_block_context || va_block_context->mm == mm);
+    mm = va_space->va_space_mm.mm;
     uvm_assert_mmap_lock_locked(mm);
     uvm_assert_rwsem_locked(&va_space->lock);
     UVM_ASSERT(PAGE_ALIGNED(addr));
 
     // Note that we have to allow PROT_NONE VMAs so that policies can be set.
-    vma = find_vma(mm, addr);
-    if (!uvm_hmm_vma_is_valid(vma, addr, allow_unreadable_vma))
+    va_block_vma = find_vma(mm, addr);
+    if (!uvm_hmm_vma_is_valid(va_block_vma, addr, allow_unreadable_vma))
         return NV_ERR_INVALID_ADDRESS;
 
     // Since we only hold the va_space read lock, there can be multiple
@@ -517,8 +531,8 @@ static NV_STATUS hmm_va_block_find_create(uvm_va_space_t *va_space,
 
 done:
     uvm_mutex_unlock(&va_space->hmm.blocks_lock);
-    if (va_block_context)
-        va_block_context->hmm.vma = vma;
+    if (vma_out)
+        *vma_out = va_block_vma;
     *va_block_ptr = va_block;
     return NV_OK;
 
@@ -532,43 +546,36 @@ err_unlock:
 
 NV_STATUS uvm_hmm_va_block_find_create(uvm_va_space_t *va_space,
                                        NvU64 addr,
-                                       uvm_va_block_context_t *va_block_context,
+                                       struct vm_area_struct **vma,
                                        uvm_va_block_t **va_block_ptr)
 {
-    return hmm_va_block_find_create(va_space, addr, false, va_block_context, va_block_ptr);
+    return hmm_va_block_find_create(va_space, addr, false, vma, va_block_ptr);
 }
 
-NV_STATUS uvm_hmm_find_vma(uvm_va_block_context_t *va_block_context, NvU64 addr)
+NV_STATUS uvm_hmm_find_vma(struct mm_struct *mm, struct vm_area_struct **vma_out, NvU64 addr)
 {
-    struct mm_struct *mm = va_block_context->mm;
-    struct vm_area_struct *vma;
-
     if (!mm)
         return NV_ERR_INVALID_ADDRESS;
 
     uvm_assert_mmap_lock_locked(mm);
 
-    vma = find_vma(mm, addr);
-    if (!uvm_hmm_vma_is_valid(vma, addr, false))
+    *vma_out = find_vma(mm, addr);
+    if (!uvm_hmm_vma_is_valid(*vma_out, addr, false))
         return NV_ERR_INVALID_ADDRESS;
 
-    va_block_context->hmm.vma = vma;
-
     return NV_OK;
 }
 
 bool uvm_hmm_check_context_vma_is_valid(uvm_va_block_t *va_block,
-                                        uvm_va_block_context_t *va_block_context,
+                                        struct vm_area_struct *vma,
                                         uvm_va_block_region_t region)
 {
     uvm_assert_mutex_locked(&va_block->lock);
 
     if (uvm_va_block_is_hmm(va_block)) {
-        struct vm_area_struct *vma = va_block_context->hmm.vma;
-
         UVM_ASSERT(vma);
-        UVM_ASSERT(va_block_context->mm == vma->vm_mm);
-        uvm_assert_mmap_lock_locked(va_block_context->mm);
+        UVM_ASSERT(va_block->hmm.va_space->va_space_mm.mm == vma->vm_mm);
+        uvm_assert_mmap_lock_locked(va_block->hmm.va_space->va_space_mm.mm);
         UVM_ASSERT(vma->vm_start <= uvm_va_block_region_start(va_block, region));
         UVM_ASSERT(vma->vm_end > uvm_va_block_region_end(va_block, region));
     }
@@ -619,8 +626,6 @@ static NV_STATUS hmm_migrate_range(uvm_va_block_t *va_block,
     uvm_mutex_lock(&va_block->lock);
 
     uvm_for_each_va_policy_in(policy, va_block, start, end, node, region) {
-        va_block_context->policy = policy;
-
         // Even though UVM_VA_BLOCK_RETRY_LOCKED() may unlock and relock the
         // va_block lock, the policy remains valid because we hold the mmap
         // lock so munmap can't remove the policy, and the va_space lock so the
@@ -670,7 +675,6 @@ void uvm_hmm_evict_va_blocks(uvm_va_space_t *va_space)
                 continue;
 
             block_context->hmm.vma = vma;
-            block_context->policy = &uvm_va_policy_default;
             uvm_hmm_va_block_migrate_locked(va_block,
                                             NULL,
                                             block_context,
@@ -1046,11 +1050,7 @@ static NV_STATUS hmm_set_preferred_location_locked(uvm_va_block_t *va_block,
             uvm_processor_mask_test(&old_policy->accessed_by, old_policy->preferred_location))
             uvm_processor_mask_set(&set_accessed_by_processors, old_policy->preferred_location);
 
-        va_block_context->policy = uvm_va_policy_set_preferred_location(va_block,
-                                                                        region,
-                                                                        preferred_location,
-                                                                        old_policy);
-        if (!va_block_context->policy)
+        if (!uvm_va_policy_set_preferred_location(va_block, region, preferred_location, old_policy))
             return NV_ERR_NO_MEMORY;
 
         // Establish new remote mappings if the old preferred location had
@@ -1109,7 +1109,7 @@ NV_STATUS uvm_hmm_set_preferred_location(uvm_va_space_t *va_space,
     for (addr = base; addr < last_address; addr = va_block->end + 1) {
         NvU64 end;
 
-        status = hmm_va_block_find_create(va_space, addr, true, va_block_context, &va_block);
+        status = hmm_va_block_find_create(va_space, addr, true, &va_block_context->hmm.vma, &va_block);
         if (status != NV_OK)
             break;
 
@@ -1151,7 +1151,6 @@ static NV_STATUS hmm_set_accessed_by_start_end_locked(uvm_va_block_t *va_block,
         if (uvm_va_policy_is_read_duplicate(&node->policy, va_space))
             continue;
 
-        va_block_context->policy = &node->policy;
         region = uvm_va_block_region_from_start_end(va_block,
                                                     max(start, node->node.start),
                                                     min(end, node->node.end));
@@ -1196,7 +1195,7 @@ NV_STATUS uvm_hmm_set_accessed_by(uvm_va_space_t *va_space,
     for (addr = base; addr < last_address; addr = va_block->end + 1) {
         NvU64 end;
 
-        status = hmm_va_block_find_create(va_space, addr, true, va_block_context, &va_block);
+        status = hmm_va_block_find_create(va_space, addr, true, &va_block_context->hmm.vma, &va_block);
         if (status != NV_OK)
             break;
 
@@ -1249,8 +1248,6 @@ void uvm_hmm_block_add_eviction_mappings(uvm_va_space_t *va_space,
     uvm_mutex_lock(&va_block->lock);
 
     uvm_for_each_va_policy_node_in(node, va_block, va_block->start, va_block->end) {
-        block_context->policy = &node->policy;
-
         for_each_id_in_mask(id, &node->policy.accessed_by) {
             status = hmm_set_accessed_by_start_end_locked(va_block,
                                                           block_context,
@@ -1309,13 +1306,13 @@ void uvm_hmm_block_add_eviction_mappings(uvm_va_space_t *va_space,
     }
 }
 
-void uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
-                             uvm_va_block_context_t *va_block_context,
-                             unsigned long addr,
-                             NvU64 *endp)
+const uvm_va_policy_t *uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
+                                               struct vm_area_struct *vma,
+                                               unsigned long addr,
+                                               NvU64 *endp)
 {
-    struct vm_area_struct *vma = va_block_context->hmm.vma;
     const uvm_va_policy_node_t *node;
+    const uvm_va_policy_t *policy;
     NvU64 end = va_block->end;
 
     uvm_assert_mmap_lock_locked(vma->vm_mm);
@@ -1326,40 +1323,45 @@ void uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
 
     node = uvm_va_policy_node_find(va_block, addr);
     if (node) {
-        va_block_context->policy = &node->policy;
+        policy = &node->policy;
         if (end > node->node.end)
             end = node->node.end;
     }
     else {
-        va_block_context->policy = &uvm_va_policy_default;
+        policy = &uvm_va_policy_default;
     }
 
     *endp = end;
+
+    return policy;
 }
 
 NV_STATUS uvm_hmm_find_policy_vma_and_outer(uvm_va_block_t *va_block,
-                                            uvm_va_block_context_t *va_block_context,
+                                            struct vm_area_struct **vma_out,
                                             uvm_page_index_t page_index,
+                                            const uvm_va_policy_t **policy,
                                             uvm_page_index_t *outerp)
 {
-    struct vm_area_struct *vma;
     unsigned long addr;
     NvU64 end;
     uvm_page_index_t outer;
+    uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
+    struct mm_struct *mm = va_space->va_space_mm.mm;
+
+    if (!mm)
+        return NV_ERR_INVALID_ADDRESS;
 
     UVM_ASSERT(uvm_va_block_is_hmm(va_block));
-    uvm_assert_mmap_lock_locked(va_block_context->mm);
+    uvm_assert_mmap_lock_locked(mm);
     uvm_assert_mutex_locked(&va_block->lock);
 
     addr = uvm_va_block_cpu_page_address(va_block, page_index);
 
-    vma = vma_lookup(va_block_context->mm, addr);
-    if (!vma || !(vma->vm_flags & VM_READ))
+    *vma_out = vma_lookup(mm, addr);
+    if (!*vma_out || !((*vma_out)->vm_flags & VM_READ))
         return NV_ERR_INVALID_ADDRESS;
 
-    va_block_context->hmm.vma = vma;
-
-    uvm_hmm_find_policy_end(va_block, va_block_context, addr, &end);
+    *policy = uvm_hmm_find_policy_end(va_block, *vma_out, addr, &end);
 
     outer = uvm_va_block_cpu_page_index(va_block, end) + 1;
     if (*outerp > outer)
@@ -1379,8 +1381,6 @@ static NV_STATUS hmm_clear_thrashing_policy(uvm_va_block_t *va_block,
     uvm_mutex_lock(&va_block->lock);
 
     uvm_for_each_va_policy_in(policy, va_block, va_block->start, va_block->end, node, region) {
-        block_context->policy = policy;
-
         // Unmap may split PTEs and require a retry. Needs to be called
         // before the pinned pages information is destroyed.
         status = UVM_VA_BLOCK_RETRY_LOCKED(va_block,
@@ -1424,11 +1424,10 @@ NV_STATUS uvm_hmm_clear_thrashing_policy(uvm_va_space_t *va_space)
 }
 
 uvm_va_block_region_t uvm_hmm_get_prefetch_region(uvm_va_block_t *va_block,
-                                                  uvm_va_block_context_t *va_block_context,
+                                                  struct vm_area_struct *vma,
+                                                  const uvm_va_policy_t *policy,
                                                   NvU64 address)
 {
-    struct vm_area_struct *vma = va_block_context->hmm.vma;
-    const uvm_va_policy_t *policy = va_block_context->policy;
     NvU64 start, end;
 
     UVM_ASSERT(uvm_va_block_is_hmm(va_block));
@@ -1457,13 +1456,11 @@ uvm_va_block_region_t uvm_hmm_get_prefetch_region(uvm_va_block_t *va_block,
 }
 
 uvm_prot_t uvm_hmm_compute_logical_prot(uvm_va_block_t *va_block,
-                                        uvm_va_block_context_t *va_block_context,
+                                        struct vm_area_struct *vma,
                                         NvU64 addr)
 {
-    struct vm_area_struct *vma = va_block_context->hmm.vma;
-
     UVM_ASSERT(uvm_va_block_is_hmm(va_block));
-    uvm_assert_mmap_lock_locked(va_block_context->mm);
+    uvm_assert_mmap_lock_locked(va_block->hmm.va_space->va_space_mm.mm);
     UVM_ASSERT(vma && addr >= vma->vm_start && addr < vma->vm_end);
 
     if (!(vma->vm_flags & VM_READ))
@@ -1867,7 +1864,7 @@ static void fill_dst_pfn(uvm_va_block_t *va_block,
 
     dpage = pfn_to_page(pfn);
     UVM_ASSERT(is_device_private_page(dpage));
-    UVM_ASSERT(dpage->pgmap->owner == &g_uvm_global);
+    UVM_ASSERT(page_pgmap(dpage)->owner == &g_uvm_global);
 
     hmm_mark_gpu_chunk_referenced(va_block, gpu, gpu_chunk);
     UVM_ASSERT(!page_count(dpage));
@@ -2284,6 +2281,39 @@ static void hmm_release_atomic_pages(uvm_va_block_t *va_block,
     }
 }
 
+static int hmm_make_device_exclusive_range(struct mm_struct *mm,
+                                           unsigned long start,
+                                           unsigned long end,
+                                           struct page **pages)
+{
+#if NV_IS_EXPORT_SYMBOL_PRESENT_make_device_exclusive
+    unsigned long addr;
+    int npages = 0;
+
+    for (addr = start; addr < end; addr += PAGE_SIZE) {
+        struct folio *folio;
+        struct page *page;
+
+        page = make_device_exclusive(mm, addr, &g_uvm_global, &folio);
+        if (IS_ERR(page)) {
+            while (npages) {
+                page = pages[--npages];
+                unlock_page(page);
+                put_page(page);
+            }
+            npages = PTR_ERR(page);
+            break;
+        }
+
+        pages[npages++] = page;
+    }
+
+    return npages;
+#else
+    return make_device_exclusive_range(mm, start, end, pages, &g_uvm_global);
+#endif
+}
+
 static NV_STATUS hmm_block_atomic_fault_locked(uvm_processor_id_t processor_id,
                                                uvm_va_block_t *va_block,
                                                uvm_va_block_retry_t *va_block_retry,
@@ -2339,11 +2369,10 @@ static NV_STATUS hmm_block_atomic_fault_locked(uvm_processor_id_t processor_id,
 
     uvm_mutex_unlock(&va_block->lock);
 
-    npages = make_device_exclusive_range(service_context->block_context.mm,
+    npages = hmm_make_device_exclusive_range(service_context->block_context.mm,
         uvm_va_block_cpu_page_address(va_block, region.first),
         uvm_va_block_cpu_page_address(va_block, region.outer - 1) + PAGE_SIZE,
-        pages + region.first,
-        &g_uvm_global);
+        pages + region.first);
 
     uvm_mutex_lock(&va_block->lock);
 
@@ -2575,7 +2604,7 @@ static NV_STATUS dmamap_src_sysmem_pages(uvm_va_block_t *va_block,
                 continue;
             }
 
-            if (PageSwapCache(src_page)) {
+            if (nv_PageSwapCache(src_page)) {
                 // TODO: Bug 4050579: Remove this when swap cached pages can be
                 // migrated.
                 if (service_context) {
@@ -2907,8 +2936,6 @@ static NV_STATUS uvm_hmm_migrate_alloc_and_copy(struct vm_area_struct *vma,
     if (status != NV_OK)
         return status;
 
-    UVM_ASSERT(!uvm_va_policy_is_read_duplicate(va_block_context->policy, va_block->hmm.va_space));
-
     status = uvm_va_block_make_resident_copy(va_block,
                                              va_block_retry,
                                              va_block_context,
@@ -3140,7 +3167,7 @@ NV_STATUS uvm_hmm_migrate_ranges(uvm_va_space_t *va_space,
     for (addr = base; addr < last_address; addr = end + 1) {
         struct vm_area_struct *vma;
 
-        status = hmm_va_block_find_create(va_space, addr, false, va_block_context, &va_block);
+        status = hmm_va_block_find_create(va_space, addr, false, &va_block_context->hmm.vma, &va_block);
         if (status != NV_OK)
             return status;
 
@@ -3232,7 +3259,6 @@ static NV_STATUS hmm_va_block_evict_chunks(uvm_va_block_t *va_block,
     uvm_for_each_va_policy_in(policy, va_block, start, end, node, region) {
         npages = uvm_va_block_region_num_pages(region);
 
-        va_block_context->policy = policy;
         if (out_accessed_by_set && uvm_processor_mask_get_count(&policy->accessed_by) > 0)
             *out_accessed_by_set = true;
 

kernel-open/nvidia-uvm/uvm_hmm.h

@@ -49,9 +49,7 @@ typedef struct
     bool uvm_hmm_is_enabled_system_wide(void);
 
     // Initialize HMM for the given the va_space.
-    // Locking: the va_space->va_space_mm.mm mmap_lock must be write locked
-    // and the va_space lock must be held in write mode.
-    NV_STATUS uvm_hmm_va_space_initialize(uvm_va_space_t *va_space);
+    void uvm_hmm_va_space_initialize(uvm_va_space_t *va_space);
 
     // Destroy any HMM state for the given the va_space.
     // Locking: va_space lock must be held in write mode.
@@ -90,31 +88,30 @@ typedef struct
     // address 'addr' or the VMA does not have at least PROT_READ permission.
     // The caller is also responsible for checking that there is no UVM
     // va_range covering the given address before calling this function.
-    // If va_block_context is not NULL, the VMA is cached in
-    // va_block_context->hmm.vma.
+    // The VMA is returned in vma_out if it's not NULL.
     // Locking: This function must be called with mm retained and locked for
     // at least read and the va_space lock at least for read.
     NV_STATUS uvm_hmm_va_block_find_create(uvm_va_space_t *va_space,
                                            NvU64 addr,
-                                           uvm_va_block_context_t *va_block_context,
+                                           struct vm_area_struct **vma_out,
                                            uvm_va_block_t **va_block_ptr);
 
-    // Find the VMA for the given address and set va_block_context->hmm.vma.
-    // Return NV_ERR_INVALID_ADDRESS if va_block_context->mm is NULL or there
-    // is no VMA associated with the address 'addr' or the VMA does not have at
-    // least PROT_READ permission.
+    // Find the VMA for the given address and return it in vma_out. Return
+    // NV_ERR_INVALID_ADDRESS if mm is NULL or there is no VMA associated with
+    // the address 'addr' or the VMA does not have at least PROT_READ
+    // permission.
     // Locking: This function must be called with mm retained and locked for
     // at least read or mm equal to NULL.
-    NV_STATUS uvm_hmm_find_vma(uvm_va_block_context_t *va_block_context, NvU64 addr);
+    NV_STATUS uvm_hmm_find_vma(struct mm_struct *mm, struct vm_area_struct **vma_out, NvU64 addr);
 
-    // If va_block is a HMM va_block, check that va_block_context->hmm.vma is
-    // not NULL and covers the given region. This always returns true and is
-    // intended to only be used with UVM_ASSERT().
+    // If va_block is a HMM va_block, check that vma is not NULL and covers the
+    // given region. This always returns true and is intended to only be used
+    // with UVM_ASSERT().
     // Locking: This function must be called with the va_block lock held and if
-    // va_block is a HMM block, va_block_context->mm must be retained and
-    // locked for at least read.
+    // va_block is a HMM block, va_space->va_space_mm.mm->mmap_lock must be
+    // retained and locked for at least read.
     bool uvm_hmm_check_context_vma_is_valid(uvm_va_block_t *va_block,
-                                            uvm_va_block_context_t *va_block_context,
+                                            struct vm_area_struct *vma,
                                             uvm_va_block_region_t region);
 
     // Initialize the HMM portion of the service_context.
@@ -225,31 +222,29 @@ typedef struct
         return NV_OK;
     }
 
-    // This function assigns va_block_context->policy to the policy covering
-    // the given address 'addr' and assigns the ending address '*endp' to the
-    // minimum of va_block->end, va_block_context->hmm.vma->vm_end - 1, and the
-    // ending address of the policy range. Note that va_block_context->hmm.vma
-    // is expected to be initialized before calling this function.
-    // Locking: This function must be called with
-    // va_block_context->hmm.vma->vm_mm retained and locked for least read and
-    // the va_block lock held.
-    void uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
-                                 uvm_va_block_context_t *va_block_context,
-                                 unsigned long addr,
-                                 NvU64 *endp);
+    // This function returns the policy covering the given address 'addr' and
+    // assigns the ending address '*endp' to the minimum of va_block->end,
+    // vma->vm_end - 1, and the ending address of the policy range. Locking:
+    // This function must be called with vma->vm_mm retained and locked for at
+    // least read and the va_block and va_space lock held.
+    const uvm_va_policy_t *uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
+                                                   struct vm_area_struct *vma,
+                                                   unsigned long addr,
+                                                   NvU64 *endp);
 
-    // This function finds the VMA for the page index 'page_index' and assigns
-    // it to va_block_context->vma, sets va_block_context->policy to the policy
-    // covering the given address, and sets the ending page range '*outerp'
-    // to the minimum of *outerp, va_block_context->hmm.vma->vm_end - 1, the
-    // ending address of the policy range, and va_block->end.
-    // Return NV_ERR_INVALID_ADDRESS if no VMA is found; otherwise, NV_OK.
-    // Locking: This function must be called with
-    // va_block_context->hmm.vma->vm_mm retained and locked for least read and
-    // the va_block lock held.
+    // This function finds the VMA for the page index 'page_index' and returns
+    // it in vma_out which must not be NULL. Returns the policy covering the
+    // given address, and sets the ending page range '*outerp' to the minimum of
+    // *outerp, vma->vm_end - 1, the ending address of the policy range, and
+    // va_block->end.
+    // Return NV_ERR_INVALID_ADDRESS if no VMA is found; otherwise sets *vma
+    // and returns NV_OK.
+    // Locking: This function must be called with mm retained and locked for at
+    // least read and the va_block and va_space lock held.
     NV_STATUS uvm_hmm_find_policy_vma_and_outer(uvm_va_block_t *va_block,
-                                                uvm_va_block_context_t *va_block_context,
+                                                struct vm_area_struct **vma,
                                                 uvm_page_index_t page_index,
+                                                const uvm_va_policy_t **policy,
                                                 uvm_page_index_t *outerp);
 
     // Clear thrashing policy information from all HMM va_blocks.
@@ -258,24 +253,21 @@ typedef struct
 
     // Return the expanded region around 'address' limited to the intersection
     // of va_block start/end, vma start/end, and policy start/end.
-    // va_block_context must not be NULL, va_block_context->hmm.vma must be
-    // valid (this is usually set by uvm_hmm_va_block_find_create()), and
-    // va_block_context->policy must be valid.
-    // Locking: the caller must hold mm->mmap_lock in at least read mode, the
-    // va_space lock must be held in at least read mode, and the va_block lock
-    // held.
+    // Locking: the caller must hold va_space->va_space_mm.mm->mmap_lock in at
+    // least read mode, the va_space lock must be held in at least read mode,
+    // and the va_block lock held.
     uvm_va_block_region_t uvm_hmm_get_prefetch_region(uvm_va_block_t *va_block,
-                                                      uvm_va_block_context_t *va_block_context,
+                                                      struct vm_area_struct *vma,
+                                                      const uvm_va_policy_t *policy,
                                                       NvU64 address);
 
     // Return the logical protection allowed of a HMM va_block for the page at
-    // the given address.
-    // va_block_context must not be NULL and va_block_context->hmm.vma must be
-    // valid (this is usually set by uvm_hmm_va_block_find_create()).
-    // Locking: the caller must hold va_block_context->mm mmap_lock in at least
-    // read mode.
+    // the given address within the vma which must be valid. This is usually
+    // obtained from uvm_hmm_va_block_find_create()).
+    // Locking: the caller must hold va_space->va_space_mm.mm mmap_lock in at
+    // least read mode.
     uvm_prot_t uvm_hmm_compute_logical_prot(uvm_va_block_t *va_block,
-                                            uvm_va_block_context_t *va_block_context,
+                                            struct vm_area_struct *vma,
                                             NvU64 addr);
 
     // This is called to service a GPU fault.
@@ -288,9 +280,9 @@ typedef struct
                                               uvm_service_block_context_t *service_context);
 
     // This is called to migrate a region within a HMM va_block.
-    // va_block_context must not be NULL and va_block_context->policy and
-    // va_block_context->hmm.vma must be valid.
-    // Locking: the va_block_context->mm must be retained, mmap_lock must be
+    // va_block_context must not be NULL and va_block_context->hmm.vma
+    // must be valid.
+    // Locking: the va_space->va_space_mm.mm must be retained, mmap_lock must be
     // locked, and the va_block lock held.
     NV_STATUS uvm_hmm_va_block_migrate_locked(uvm_va_block_t *va_block,
                                               uvm_va_block_retry_t *va_block_retry,
@@ -303,7 +295,7 @@ typedef struct
     // UvmMigrate().
     //
     // va_block_context must not be NULL. The caller is not required to set
-    // va_block_context->policy or va_block_context->hmm.vma.
+    // va_block_context->hmm.vma.
     //
     // Locking: the va_space->va_space_mm.mm mmap_lock must be locked and
     // the va_space read lock must be held.
@@ -412,9 +404,8 @@ typedef struct
         return false;
     }
 
-    static NV_STATUS uvm_hmm_va_space_initialize(uvm_va_space_t *va_space)
+    static void uvm_hmm_va_space_initialize(uvm_va_space_t *va_space)
     {
-        return NV_OK;
     }
 
     static void uvm_hmm_va_space_destroy(uvm_va_space_t *va_space)
@@ -440,19 +431,19 @@ typedef struct
 
     static NV_STATUS uvm_hmm_va_block_find_create(uvm_va_space_t *va_space,
                                                   NvU64 addr,
-                                                  uvm_va_block_context_t *va_block_context,
+                                                  struct vm_area_struct **vma,
                                                   uvm_va_block_t **va_block_ptr)
     {
         return NV_ERR_INVALID_ADDRESS;
     }
 
-    static NV_STATUS uvm_hmm_find_vma(uvm_va_block_context_t *va_block_context, NvU64 addr)
+    static NV_STATUS uvm_hmm_find_vma(struct mm_struct *mm, struct vm_area_struct **vma, NvU64 addr)
     {
         return NV_OK;
     }
 
     static bool uvm_hmm_check_context_vma_is_valid(uvm_va_block_t *va_block,
-                                                   uvm_va_block_context_t *va_block_context,
+                                                   struct vm_area_struct *vma,
                                                    uvm_va_block_region_t region)
     {
         return true;
@@ -533,16 +524,19 @@ typedef struct
         return NV_ERR_INVALID_ADDRESS;
     }
 
-    static void uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
-                                        uvm_va_block_context_t *va_block_context,
-                                        unsigned long addr,
-                                        NvU64 *endp)
+    static const uvm_va_policy_t *uvm_hmm_find_policy_end(uvm_va_block_t *va_block,
+                                                          struct vm_area_struct *vma,
+                                                          unsigned long addr,
+                                                          NvU64 *endp)
     {
+        UVM_ASSERT(0);
+        return NULL;
     }
 
     static NV_STATUS uvm_hmm_find_policy_vma_and_outer(uvm_va_block_t *va_block,
-                                                       uvm_va_block_context_t *va_block_context,
+                                                       struct vm_area_struct **vma,
                                                        uvm_page_index_t page_index,
+                                                       const uvm_va_policy_t **policy,
                                                        uvm_page_index_t *outerp)
     {
         return NV_OK;
@@ -554,14 +548,15 @@ typedef struct
     }
 
     static uvm_va_block_region_t uvm_hmm_get_prefetch_region(uvm_va_block_t *va_block,
-                                                             uvm_va_block_context_t *va_block_context,
+                                                             struct vm_area_struct *vma,
+                                                             const uvm_va_policy_t *policy,
                                                              NvU64 address)
     {
         return (uvm_va_block_region_t){};
     }
 
     static uvm_prot_t uvm_hmm_compute_logical_prot(uvm_va_block_t *va_block,
-                                                   uvm_va_block_context_t *va_block_context,
+                                                   struct vm_area_struct *vma,
                                                    NvU64 addr)
     {
         return UVM_PROT_NONE;

kernel-open/nvidia-uvm/uvm_hopper.c

@@ -61,7 +61,11 @@ void uvm_hal_hopper_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     // GH180.
     parent_gpu->ce_phys_vidmem_write_supported = !uvm_gpu_is_coherent(parent_gpu);
 
-    parent_gpu->peer_copy_mode = g_uvm_global.peer_copy_mode;
+    // TODO: Bug 4174553: [HGX-SkinnyJoe][GH180] channel errors discussion/debug
+    //                    portion for the uvm tests became nonresponsive after
+    //                    some time and then failed even after reboot
+    parent_gpu->peer_copy_mode = uvm_gpu_is_coherent(parent_gpu) ?
+                                                           UVM_GPU_PEER_COPY_MODE_VIRTUAL : g_uvm_global.peer_copy_mode;
 
     // All GR context buffers may be mapped to 57b wide VAs. All "compute" units
     // accessing GR context buffers support the 57-bit VA range.
@@ -99,5 +103,7 @@ void uvm_hal_hopper_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     parent_gpu->map_remap_larger_page_promotion = false;
 
     parent_gpu->plc_supported = true;
+
+    parent_gpu->no_ats_range_required = true;
 }
 

kernel-open/nvidia-uvm/uvm_hopper_ce.c

@@ -491,7 +491,6 @@ void uvm_hal_hopper_ce_encrypt(uvm_push_t *push,
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
 
     UVM_ASSERT(uvm_conf_computing_mode_is_hcc(gpu));
-    UVM_ASSERT(uvm_push_is_fake(push) || uvm_channel_is_secure(push->channel));
     UVM_ASSERT(IS_ALIGNED(auth_tag.address, UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT));
 
     if (!src.is_virtual)
@@ -540,7 +539,6 @@ void uvm_hal_hopper_ce_decrypt(uvm_push_t *push,
     uvm_gpu_t *gpu = uvm_push_get_gpu(push);
 
     UVM_ASSERT(uvm_conf_computing_mode_is_hcc(gpu));
-    UVM_ASSERT(!push->channel || uvm_channel_is_secure(push->channel));
     UVM_ASSERT(IS_ALIGNED(auth_tag.address, UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT));
 
     // The addressing mode (and aperture, if applicable) of the source and

kernel-open/nvidia-uvm/uvm_hopper_mmu.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2020-2022 NVIDIA Corporation
+    Copyright (c) 2020-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -33,6 +33,7 @@
 
 #include "uvm_types.h"
 #include "uvm_global.h"
+#include "uvm_common.h"
 #include "uvm_hal.h"
 #include "uvm_hal_types.h"
 #include "uvm_hopper_fault_buffer.h"
@@ -42,6 +43,10 @@
 #define MMU_BIG 0
 #define MMU_SMALL 1
 
+// Used in pde_pcf().
+#define ATS_ALLOWED 0
+#define ATS_NOT_ALLOWED 1
+
 uvm_mmu_engine_type_t uvm_hal_hopper_mmu_engine_id_to_type(NvU16 mmu_engine_id)
 {
     if (mmu_engine_id >= NV_PFAULT_MMU_ENG_ID_HOST0 && mmu_engine_id <= NV_PFAULT_MMU_ENG_ID_HOST44)
@@ -260,7 +265,108 @@ static NvU64 poisoned_pte_hopper(void)
     return WRITE_HWCONST64(pte_bits, _MMU_VER3, PTE, PCF, PRIVILEGE_RO_NO_ATOMIC_UNCACHED_ACD);
 }
 
-static NvU64 single_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc, NvU32 depth)
+typedef enum
+{
+    PDE_TYPE_SINGLE,
+    PDE_TYPE_DUAL_BIG,
+    PDE_TYPE_DUAL_SMALL,
+    PDE_TYPE_COUNT,
+} pde_type_t;
+
+static const NvU8 valid_pcf[][2] = { { NV_MMU_VER3_PDE_PCF_VALID_UNCACHED_ATS_ALLOWED,
+                                       NV_MMU_VER3_PDE_PCF_VALID_UNCACHED_ATS_NOT_ALLOWED },
+                                     { NV_MMU_VER3_DUAL_PDE_PCF_BIG_VALID_UNCACHED_ATS_ALLOWED,
+                                       NV_MMU_VER3_DUAL_PDE_PCF_BIG_VALID_UNCACHED_ATS_NOT_ALLOWED },
+                                     { NV_MMU_VER3_DUAL_PDE_PCF_SMALL_VALID_UNCACHED_ATS_ALLOWED,
+                                       NV_MMU_VER3_DUAL_PDE_PCF_SMALL_VALID_UNCACHED_ATS_NOT_ALLOWED } };
+
+static const NvU8 invalid_pcf[][2] = { { NV_MMU_VER3_PDE_PCF_INVALID_ATS_ALLOWED,
+                                         NV_MMU_VER3_PDE_PCF_INVALID_ATS_NOT_ALLOWED },
+                                       { NV_MMU_VER3_DUAL_PDE_PCF_BIG_INVALID_ATS_ALLOWED,
+                                         NV_MMU_VER3_DUAL_PDE_PCF_BIG_INVALID_ATS_NOT_ALLOWED },
+                                       { NV_MMU_VER3_DUAL_PDE_PCF_SMALL_INVALID_ATS_ALLOWED,
+                                         NV_MMU_VER3_DUAL_PDE_PCF_SMALL_INVALID_ATS_NOT_ALLOWED } };
+
+static const NvU8 va_base[] = { 56, 47, 38, 29, 21 };
+
+static bool is_ats_range_valid(uvm_page_directory_t *dir, NvU32 child_index)
+{
+    NvU64 pde_base_va;
+    NvU64 min_va_upper;
+    NvU64 max_va_lower;
+    NvU32 index_in_dir;
+
+    uvm_cpu_get_unaddressable_range(&max_va_lower, &min_va_upper);
+
+    UVM_ASSERT(dir->depth < ARRAY_SIZE(va_base));
+
+    // We can use UVM_PAGE_SIZE_AGNOSTIC because page_size is only used in
+    // index_bits_hopper() for PTE table, i.e., depth 5+, which does not use a
+    // PDE PCF or an ATS_ALLOWED/NOT_ALLOWED setting.
+    UVM_ASSERT(child_index < (1ull << index_bits_hopper(dir->depth, UVM_PAGE_SIZE_AGNOSTIC)));
+
+    pde_base_va = 0;
+    index_in_dir = child_index;
+    while (dir) {
+        pde_base_va += index_in_dir * (1ull << va_base[dir->depth]);
+        index_in_dir = dir->index_in_parent;
+        dir = dir->host_parent;
+    }
+    pde_base_va = (NvU64)((NvS64)(pde_base_va << (64 - num_va_bits_hopper())) >> (64 - num_va_bits_hopper()));
+
+    if (pde_base_va < max_va_lower || pde_base_va >= min_va_upper)
+        return true;
+
+    return false;
+}
+
+// PDE Permission Control Flags
+static NvU32 pde_pcf(bool valid, pde_type_t pde_type, uvm_page_directory_t *dir, NvU32 child_index)
+{
+    const NvU8 (*pcf)[2] = valid ? valid_pcf : invalid_pcf;
+    NvU8 depth = dir->depth;
+
+    UVM_ASSERT(pde_type < PDE_TYPE_COUNT);
+    UVM_ASSERT(depth < 5);
+
+    // On non-ATS systems, PDE PCF only sets the valid and volatile/cache bits.
+    if (!g_uvm_global.ats.enabled)
+        return pcf[pde_type][ATS_ALLOWED];
+
+    // We assume all supported ATS platforms use canonical form address.
+    // See comments in uvm_gpu.c:uvm_gpu_can_address() and in
+    // uvm_mmu.c:page_tree_ats_init();
+    UVM_ASSERT(uvm_platform_uses_canonical_form_address());
+
+    // Hopper GPUs on ATS-enabled systems, perform a parallel lookup on both
+    // ATS and GMMU page tables. For managed memory we need to prevent this
+    // parallel lookup since we would not get any GPU fault if the CPU has
+    // a valid mapping. Also, for external ranges that are known to be
+    // mapped entirely on the GMMU page table we can skip the ATS lookup
+    // for performance reasons. Parallel ATS lookup is disabled in PDE1
+    // (depth 3) and, therefore, it applies to the underlying 512MB VA
+    // range.
+    //
+    // UVM sets ATS_NOT_ALLOWED for all Hopper+ mappings on ATS systems.
+    // This is fine because CUDA ensures that all managed and external
+    // allocations are properly compartmentalized in 512MB-aligned VA
+    // regions. For cudaHostRegister CUDA cannot control the VA range, but
+    // we rely on ATS for those allocations so they can't choose the
+    // ATS_NOT_ALLOWED mode.
+    // TODO: Bug 3254055: Relax the NO_ATS setting from 512MB (pde1) range to
+    //                    PTEs.
+    // HW complies with the leaf PDE's ATS_ALLOWED/ATS_NOT_ALLOWED settings,
+    // enabling us to treat any upper-level PDE as a don't care as long as there
+    // are leaf PDEs for the entire upper-level PDE range. We assume PDE4
+    // entries (depth == 0) are always ATS enabled, and the no_ats_range is in
+    // PDE3 or lower.
+    if (depth == 0 || (!valid && is_ats_range_valid(dir, child_index)))
+        return pcf[pde_type][ATS_ALLOWED];
+
+    return pcf[pde_type][ATS_NOT_ALLOWED];
+}
+
+static NvU64 single_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc, uvm_page_directory_t *dir, NvU32 child_index)
 {
     NvU64 pde_bits = 0;
 
@@ -280,38 +386,17 @@ static NvU64 single_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc, NvU32 dep
                 break;
         }
 
-        // PCF (permission control flags) 5:3
-        // Hopper GPUs on ATS-enabled systems, perform a parallel lookup on both
-        // ATS and GMMU page tables. For managed memory we need to prevent this
-        // parallel lookup since we would not get any GPU fault if the CPU has
-        // a valid mapping. Also, for external ranges that are known to be
-        // mapped entirely on the GMMU page table we can skip the ATS lookup
-        // for performance reasons. Parallel ATS lookup is disabled in PDE1
-        // (depth 3) and, therefore, it applies to the underlying 512MB VA
-        // range.
-        //
-        // UVM sets ATS_NOT_ALLOWED for all Hopper+ mappings on ATS systems.
-        // This is fine because CUDA ensures that all managed and external
-        // allocations are properly compartmentalized in 512MB-aligned VA
-        // regions. For cudaHostRegister CUDA cannot control the VA range, but
-        // we rely on ATS for those allocations so they can't choose the
-        // ATS_NOT_ALLOWED mode.
-        //
-        // TODO: Bug 3254055: Relax the NO_ATS setting from 512MB (pde1) range
-        // to PTEs.
-        if (depth == 3 && g_uvm_global.ats.enabled)
-            pde_bits |= HWCONST64(_MMU_VER3, PDE, PCF, VALID_UNCACHED_ATS_NOT_ALLOWED);
-        else
-            pde_bits |= HWCONST64(_MMU_VER3, PDE, PCF, VALID_UNCACHED_ATS_ALLOWED);
-
         // address 51:12
         pde_bits |= HWVALUE64(_MMU_VER3, PDE, ADDRESS, address);
     }
 
+    // PCF (permission control flags) 5:3
+    pde_bits |= HWVALUE64(_MMU_VER3, PDE, PCF, pde_pcf(phys_alloc != NULL, PDE_TYPE_SINGLE, dir, child_index));
+
     return pde_bits;
 }
 
-static NvU64 big_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc)
+static NvU64 big_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc, uvm_page_directory_t *dir, NvU32 child_index)
 {
     NvU64 pde_bits = 0;
 
@@ -330,17 +415,20 @@ static NvU64 big_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc)
                 break;
         }
 
-        // PCF (permission control flags) 5:3
-        pde_bits |= HWCONST64(_MMU_VER3, DUAL_PDE, PCF_BIG, VALID_UNCACHED_ATS_NOT_ALLOWED);
-
         // address 51:8
         pde_bits |= HWVALUE64(_MMU_VER3, DUAL_PDE, ADDRESS_BIG, address);
     }
 
+    // PCF (permission control flags) 5:3
+    pde_bits |= HWVALUE64(_MMU_VER3,
+                          DUAL_PDE,
+                          PCF_BIG,
+                          pde_pcf(phys_alloc != NULL, PDE_TYPE_DUAL_BIG, dir, child_index));
+
     return pde_bits;
 }
 
-static NvU64 small_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc)
+static NvU64 small_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc, uvm_page_directory_t *dir, NvU32 child_index)
 {
     NvU64 pde_bits = 0;
 
@@ -359,29 +447,40 @@ static NvU64 small_half_pde_hopper(uvm_mmu_page_table_alloc_t *phys_alloc)
                 break;
         }
 
-        // PCF (permission control flags) 69:67 [5:3]
-        pde_bits |= HWCONST64(_MMU_VER3, DUAL_PDE, PCF_SMALL, VALID_UNCACHED_ATS_NOT_ALLOWED);
-
         // address 115:76 [51:12]
         pde_bits |= HWVALUE64(_MMU_VER3, DUAL_PDE, ADDRESS_SMALL, address);
     }
+
+    // PCF (permission control flags) 69:67 [5:3]
+    pde_bits |= HWVALUE64(_MMU_VER3,
+                          DUAL_PDE,
+                          PCF_SMALL,
+                          pde_pcf(phys_alloc != NULL, PDE_TYPE_DUAL_SMALL, dir, child_index));
+
     return pde_bits;
 }
 
-static void make_pde_hopper(void *entry, uvm_mmu_page_table_alloc_t **phys_allocs, NvU32 depth)
+static void make_pde_hopper(void *entry,
+                            uvm_mmu_page_table_alloc_t **phys_allocs,
+                            uvm_page_directory_t *dir,
+                            NvU32 child_index)
 {
-    NvU32 entry_count = entries_per_index_hopper(depth);
+    NvU32 entry_count;
     NvU64 *entry_bits = (NvU64 *)entry;
 
+    UVM_ASSERT(dir);
+
+    entry_count = entries_per_index_hopper(dir->depth);
+
     if (entry_count == 1) {
-        *entry_bits = single_pde_hopper(*phys_allocs, depth);
+        *entry_bits = single_pde_hopper(*phys_allocs, dir, child_index);
     }
     else if (entry_count == 2) {
-        entry_bits[MMU_BIG] = big_half_pde_hopper(phys_allocs[MMU_BIG]);
-        entry_bits[MMU_SMALL] = small_half_pde_hopper(phys_allocs[MMU_SMALL]);
+        entry_bits[MMU_BIG] = big_half_pde_hopper(phys_allocs[MMU_BIG], dir, child_index);
+        entry_bits[MMU_SMALL] = small_half_pde_hopper(phys_allocs[MMU_SMALL], dir, child_index);
 
         // This entry applies to the whole dual PDE but is stored in the lower
-        // bits
+        // bits.
         entry_bits[MMU_BIG] |= HWCONST64(_MMU_VER3, DUAL_PDE, IS_PTE, FALSE);
     }
     else {

kernel-open/nvidia-uvm/uvm_kvmalloc.c

@@ -36,7 +36,7 @@
 typedef struct
 {
     size_t alloc_size;
-    uint8_t ptr[0];
+    uint8_t ptr[];
 } uvm_vmalloc_hdr_t;
 
 typedef struct

kernel-open/nvidia-uvm/uvm_linux.h

@@ -114,6 +114,16 @@ static inline const struct cpumask *uvm_cpumask_of_node(int node)
         #define UVM_IS_CONFIG_HMM() 0
     #endif
 
+// ATS prefetcher uses hmm_range_fault() to query residency information.
+// hmm_range_fault() needs CONFIG_HMM_MIRROR. To detect racing CPU invalidates
+// of memory regions while hmm_range_fault() is being called, MMU interval
+// notifiers are needed.
+    #if defined(CONFIG_HMM_MIRROR) && defined(NV_MMU_INTERVAL_NOTIFIER)
+        #define UVM_HMM_RANGE_FAULT_SUPPORTED() 1
+    #else
+        #define UVM_HMM_RANGE_FAULT_SUPPORTED() 0
+    #endif
+
 // Various issues prevent us from using mmu_notifiers in older kernels. These
 // include:
 //  - ->release being called under RCU instead of SRCU: fixed by commit
@@ -128,8 +138,9 @@ static inline const struct cpumask *uvm_cpumask_of_node(int node)
 // present if we see the callback.
 //
 // The callback was added in commit 0f0a327fa12cd55de5e7f8c05a70ac3d047f405e,
-// v3.19 (2014-11-13).
-    #if defined(NV_MMU_NOTIFIER_OPS_HAS_INVALIDATE_RANGE)
+// v3.19 (2014-11-13) and renamed in commit 1af5a8109904.
+    #if defined(NV_MMU_NOTIFIER_OPS_HAS_INVALIDATE_RANGE) || \
+        defined(NV_MMU_NOTIFIER_OPS_HAS_ARCH_INVALIDATE_SECONDARY_TLBS)
         #define UVM_CAN_USE_MMU_NOTIFIERS() 1
     #else
         #define UVM_CAN_USE_MMU_NOTIFIERS() 0
@@ -153,10 +164,6 @@ static inline const struct cpumask *uvm_cpumask_of_node(int node)
 #define VM_MIXEDMAP    0x00000000
 #endif
 
-#if !defined(MPOL_PREFERRED_MANY)
-#define MPOL_PREFERRED_MANY    5
-#endif
-
 //
 // printk.h already defined pr_fmt, so we have to redefine it so the pr_*
 // routines pick up our version
@@ -219,7 +226,7 @@ static inline const struct cpumask *uvm_cpumask_of_node(int node)
 #define __GFP_NORETRY 0
 #endif
 
-#define NV_UVM_GFP_FLAGS (GFP_KERNEL)
+#define NV_UVM_GFP_FLAGS (GFP_KERNEL | __GFP_NOMEMALLOC)
 
 // Develop builds define DEBUG but enable optimization
 #if defined(DEBUG) && !defined(NVIDIA_UVM_DEVELOP)
@@ -583,4 +590,9 @@ static inline pgprot_t uvm_pgprot_decrypted(pgprot_t prot)
   #include <asm/page.h>
   #define page_to_virt(x)    __va(PFN_PHYS(page_to_pfn(x)))
 #endif
+
+#ifndef NV_PAGE_PGMAP_PRESENT
+#define page_pgmap(page) (page)->pgmap
+#endif
+
 #endif // _UVM_LINUX_H

kernel-open/nvidia-uvm/uvm_lock.h

@@ -279,13 +279,14 @@
 //      Operations not allowed while holding the lock:
 //      - GPU memory allocation which can evict memory (would require nesting
 //        block locks)
+//
 // - GPU DMA Allocation pool lock (gpu->conf_computing.dma_buffer_pool.lock)
 //      Order: UVM_LOCK_ORDER_CONF_COMPUTING_DMA_BUFFER_POOL
+//      Condition: The Confidential Computing feature is enabled
 //      Exclusive lock (mutex)
 //
 //      Protects:
 //      - Protect the state of the uvm_conf_computing_dma_buffer_pool_t
-//        when the Confidential Computing feature is enabled on the system.
 //
 // - Chunk mapping lock (gpu->root_chunk_mappings.bitlocks and
 //   gpu->sysmem_mappings.bitlock)
@@ -321,22 +322,25 @@
 //      Operations not allowed while holding this lock
 //      - GPU memory allocation which can evict
 //
-// - Secure channel CSL channel pool semaphore
+// - CE channel CSL channel pool semaphore
 //      Order: UVM_LOCK_ORDER_CSL_PUSH
-//      Semaphore per SEC2 channel pool
+//      Condition: The Confidential Computing feature is enabled
+//      Semaphore per CE channel pool
 //
-//      The semaphore controls concurrent pushes to secure channels. Secure work
-//      submission depends on channel availability in GPFIFO entries (as in any
-//      other channel type) but also on channel locking. Each secure channel has a
-//      lock to enforce ordering of pushes. The channel's CSL lock is taken on
-//      channel reservation until uvm_push_end. Secure channels are stateful
-//      channels and the CSL lock protects their CSL state/context.
+//      The semaphore controls concurrent pushes to CE channels that are not WCL
+//      channels. Secure work submission depends on channel availability in
+//      GPFIFO entries (as in any other channel type) but also on channel
+//      locking. Each channel has a lock to enforce ordering of pushes. The
+//      channel's CSL lock is taken on channel reservation until uvm_push_end.
+//      When the Confidential Computing feature is enabled, channels are
+//      stateful, and the CSL lock protects their CSL state/context.
 //
 //      Operations allowed while holding this lock
-//      - Pushing work to CE secure channels
+//      - Pushing work to CE channels (except for WLC channels)
 //
 // - WLC CSL channel pool semaphore
 //      Order: UVM_LOCK_ORDER_CSL_WLC_PUSH
+//      Condition: The Confidential Computing feature is enabled
 //      Semaphore per WLC channel pool
 //
 //      The semaphore controls concurrent pushes to WLC channels. WLC work
@@ -346,8 +350,8 @@
 //      channel reservation until uvm_push_end. SEC2 channels are stateful
 //      channels and the CSL lock protects their CSL state/context.
 //
-//      This lock ORDER is different and sits below generic secure channel CSL
-//      lock and above SEC2 CSL lock. This reflects the dual nature of WLC
+//      This lock ORDER is different and sits below the generic channel CSL
+//      lock and above the SEC2 CSL lock. This reflects the dual nature of WLC
 //      channels; they use SEC2 indirect work launch during initialization,
 //      and after their schedule is initialized they provide indirect launch
 //      functionality to other CE channels.
@@ -357,6 +361,7 @@
 //
 // - SEC2 CSL channel pool semaphore
 //      Order: UVM_LOCK_ORDER_SEC2_CSL_PUSH
+//      Condition: The Confidential Computing feature is enabled
 //      Semaphore per SEC2 channel pool
 //
 //      The semaphore controls concurrent pushes to SEC2 channels. SEC2 work
@@ -366,9 +371,9 @@
 //      channel reservation until uvm_push_end. SEC2 channels are stateful
 //      channels and the CSL lock protects their CSL state/context.
 //
-//      This lock ORDER is different and lower than the generic secure channel
-//      lock to allow secure work submission to use a SEC2 channel to submit
-//      work before releasing the CSL lock of the originating secure channel.
+//      This lock ORDER is different and lower than UVM_LOCK_ORDER_CSL_PUSH
+//      to allow secure work submission to use a SEC2 channel to submit
+//      work before releasing the CSL lock of the originating channel.
 //
 //      Operations allowed while holding this lock
 //      - Pushing work to SEC2 channels
@@ -408,16 +413,18 @@
 //
 // - WLC Channel lock
 //      Order: UVM_LOCK_ORDER_WLC_CHANNEL
+//      Condition: The Confidential Computing feature is enabled
 //      Spinlock (uvm_spinlock_t)
 //
 //      Lock protecting the state of WLC channels in a channel pool. This lock
-//      is separate from the above generic channel lock to allow for indirect
-//      worklaunch pushes while holding the main channel lock.
-//      (WLC pushes don't need any of the pushbuffer locks described above)
+//      is separate from the generic channel lock (UVM_LOCK_ORDER_CHANNEL)
+//      to allow for indirect worklaunch pushes while holding the main channel
+//      lock (WLC pushes don't need any of the pushbuffer locks described
+//      above)
 //
 // - Tools global VA space list lock (g_tools_va_space_list_lock)
 //      Order: UVM_LOCK_ORDER_TOOLS_VA_SPACE_LIST
-//      Reader/writer lock (rw_sempahore)
+//      Reader/writer lock (rw_semaphore)
 //
 //      This lock protects the list of VA spaces used when broadcasting
 //      UVM profiling events.
@@ -437,9 +444,10 @@
 //
 // - Tracking semaphores
 //      Order: UVM_LOCK_ORDER_SECURE_SEMAPHORE
-//      When the Confidential Computing feature is enabled, CE semaphores are
-//      encrypted, and require to take the CSL lock (UVM_LOCK_ORDER_LEAF) to
-//      decrypt the payload.
+//      Condition: The Confidential Computing feature is enabled
+//
+//      CE semaphore payloads are encrypted, and require to take the CSL lock
+//      (UVM_LOCK_ORDER_LEAF) to decrypt the payload.
 //
 // - Leaf locks
 //      Order: UVM_LOCK_ORDER_LEAF

kernel-open/nvidia-uvm/uvm_maxwell.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016-2021 NVIDIA Corporation
+    Copyright (c) 2016-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -71,4 +71,6 @@ void uvm_hal_maxwell_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     parent_gpu->smc.supported = false;
 
     parent_gpu->plc_supported = false;
+
+    parent_gpu->no_ats_range_required = false;
 }

kernel-open/nvidia-uvm/uvm_maxwell_mmu.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016-2021 NVIDIA Corporation
+    Copyright (c) 2016-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -106,10 +106,16 @@ static NvU64 small_half_pde_maxwell(uvm_mmu_page_table_alloc_t *phys_alloc)
     return pde_bits;
 }
 
-static void make_pde_maxwell(void *entry, uvm_mmu_page_table_alloc_t **phys_allocs, NvU32 depth)
+static void make_pde_maxwell(void *entry,
+                             uvm_mmu_page_table_alloc_t **phys_allocs,
+                             uvm_page_directory_t *dir,
+                             NvU32 child_index)
 {
     NvU64 pde_bits = 0;
-    UVM_ASSERT(depth == 0);
+
+    UVM_ASSERT(dir);
+    UVM_ASSERT(dir->depth == 0);
+
     pde_bits |= HWCONST64(_MMU, PDE, SIZE, FULL);
     pde_bits |= big_half_pde_maxwell(phys_allocs[MMU_BIG]) | small_half_pde_maxwell(phys_allocs[MMU_SMALL]);
 

kernel-open/nvidia-uvm/uvm_mem.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016-2022 NVIDIA Corporation
+    Copyright (c) 2016-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -93,8 +93,9 @@ static bool sysmem_can_be_mapped_on_gpu(uvm_mem_t *sysmem)
 {
     UVM_ASSERT(uvm_mem_is_sysmem(sysmem));
 
-    // If SEV is enabled, only unprotected memory can be mapped
-    if (g_uvm_global.sev_enabled)
+    // In Confidential Computing, only unprotected memory can be mapped on the
+    // GPU
+    if (g_uvm_global.conf_computing_enabled)
         return uvm_mem_is_sysmem_dma(sysmem);
 
     return true;
@@ -737,7 +738,7 @@ static NV_STATUS mem_map_cpu_to_sysmem_kernel(uvm_mem_t *mem)
             pages[page_index] = mem_cpu_page(mem, page_index * PAGE_SIZE);
     }
 
-    if (g_uvm_global.sev_enabled && uvm_mem_is_sysmem_dma(mem))
+    if (g_uvm_global.conf_computing_enabled && uvm_mem_is_sysmem_dma(mem))
         prot = uvm_pgprot_decrypted(PAGE_KERNEL_NOENC);
 
     mem->kernel.cpu_addr = vmap(pages, num_pages, VM_MAP, prot);

kernel-open/nvidia-uvm/uvm_mem.h

@@ -392,12 +392,6 @@ static NV_STATUS uvm_mem_alloc_vidmem(NvU64 size, uvm_gpu_t *gpu, uvm_mem_t **me
     return uvm_mem_alloc(&params, mem_out);
 }
 
-// Helper for allocating protected vidmem with the default page size
-static NV_STATUS uvm_mem_alloc_vidmem_protected(NvU64 size, uvm_gpu_t *gpu, uvm_mem_t **mem_out)
-{
-    return uvm_mem_alloc_vidmem(size, gpu, mem_out);
-}
-
 // Helper for allocating sysmem and mapping it on the CPU
 static NV_STATUS uvm_mem_alloc_sysmem_and_map_cpu_kernel(NvU64 size, struct mm_struct *mm, uvm_mem_t **mem_out)
 {

kernel-open/nvidia-uvm/uvm_mem_test.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016-2021 NVIDIA Corporation
+    Copyright (c) 2016-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -44,10 +44,10 @@ static NvU32 first_page_size(NvU32 page_sizes)
 
 static inline NV_STATUS __alloc_map_sysmem(NvU64 size, uvm_gpu_t *gpu, uvm_mem_t **sys_mem)
 {
-    if (g_uvm_global.sev_enabled)
+    if (g_uvm_global.conf_computing_enabled)
         return uvm_mem_alloc_sysmem_dma_and_map_cpu_kernel(size, gpu, current->mm, sys_mem);
-    else
-        return uvm_mem_alloc_sysmem_and_map_cpu_kernel(size, current->mm, sys_mem);
+
+    return uvm_mem_alloc_sysmem_and_map_cpu_kernel(size, current->mm, sys_mem);
 }
 
 static NV_STATUS check_accessible_from_gpu(uvm_gpu_t *gpu, uvm_mem_t *mem)
@@ -335,9 +335,6 @@ error:
 
 static bool should_test_page_size(size_t alloc_size, NvU32 page_size)
 {
-    if (g_uvm_global.sev_enabled)
-        return false;
-
     if (g_uvm_global.num_simulated_devices == 0)
         return true;
 

kernel-open/nvidia-uvm/uvm_migrate.c

@@ -134,6 +134,22 @@ static NV_STATUS block_migrate_map_unmapped_pages(uvm_va_block_t *va_block,
     // first map operation
     uvm_page_mask_complement(&va_block_context->caller_page_mask, &va_block->maybe_mapped_pages);
 
+    if (uvm_va_block_is_hmm(va_block) && !UVM_ID_IS_CPU(dest_id)) {
+        // Do not map pages that are already resident on the CPU. This is in
+        // order to avoid breaking system-wide atomic operations on HMM. HMM's
+        // implementation of system-side atomic operations involves restricting
+        // mappings to one processor (CPU or a GPU) at a time. If we were to
+        // grant a GPU a mapping to system memory, this gets into trouble
+        // because, on the CPU side, Linux can silently upgrade PTE permissions
+        // (move from read-only, to read-write, without any MMU notifiers
+        // firing), thus breaking the model by allowing simultaneous read-write
+        // access from two separate processors. To avoid that, just don't map
+        // such pages at all, when migrating.
+        uvm_page_mask_andnot(&va_block_context->caller_page_mask,
+                             &va_block_context->caller_page_mask,
+                             uvm_va_block_resident_mask_get(va_block, UVM_ID_CPU));
+    }
+
     // Only map those pages that are not mapped anywhere else (likely due
     // to a first touch or a migration). We pass
     // UvmEventMapRemoteCauseInvalid since the destination processor of a
@@ -207,7 +223,7 @@ NV_STATUS uvm_va_block_migrate_locked(uvm_va_block_t *va_block,
     NV_STATUS status, tracker_status = NV_OK;
 
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block, va_block_context, region));
+    UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block, va_block_context->hmm.vma, region));
 
     if (uvm_va_block_is_hmm(va_block)) {
         status = uvm_hmm_va_block_migrate_locked(va_block,
@@ -218,9 +234,9 @@ NV_STATUS uvm_va_block_migrate_locked(uvm_va_block_t *va_block,
                                                  UVM_MAKE_RESIDENT_CAUSE_API_MIGRATE);
     }
     else {
-        va_block_context->policy = uvm_va_range_get_policy(va_block->va_range);
+        uvm_va_policy_t *policy = uvm_va_range_get_policy(va_block->va_range);
 
-        if (uvm_va_policy_is_read_duplicate(va_block_context->policy, va_space)) {
+        if (uvm_va_policy_is_read_duplicate(policy, va_space)) {
             status = uvm_va_block_make_resident_read_duplicate(va_block,
                                                                va_block_retry,
                                                                va_block_context,
@@ -355,8 +371,6 @@ static bool va_block_should_do_cpu_preunmap(uvm_va_block_t *va_block,
     if (!va_block)
         return true;
 
-    UVM_ASSERT(va_range_should_do_cpu_preunmap(va_block_context->policy, uvm_va_block_get_va_space(va_block)));
-
     region = uvm_va_block_region_from_start_end(va_block, max(start, va_block->start), min(end, va_block->end));
 
     uvm_mutex_lock(&va_block->lock);
@@ -480,11 +494,9 @@ static NV_STATUS uvm_va_range_migrate(uvm_va_range_t *va_range,
                                       uvm_tracker_t *out_tracker)
 {
     NvU64 preunmap_range_start = start;
+    uvm_va_policy_t *policy = uvm_va_range_get_policy(va_range);
 
-    UVM_ASSERT(va_block_context->policy == uvm_va_range_get_policy(va_range));
-
-    should_do_cpu_preunmap = should_do_cpu_preunmap && va_range_should_do_cpu_preunmap(va_block_context->policy,
-                                                                                       va_range->va_space);
+    should_do_cpu_preunmap = should_do_cpu_preunmap && va_range_should_do_cpu_preunmap(policy, va_range->va_space);
 
     // Divide migrations into groups of contiguous VA blocks. This is to trigger
     // CPU unmaps for that region before the migration starts.
@@ -561,8 +573,6 @@ static NV_STATUS uvm_migrate_ranges(uvm_va_space_t *va_space,
             break;
         }
 
-        va_block_context->policy = uvm_va_range_get_policy(va_range);
-
         // For UVM-Lite GPUs, the CUDA driver may suballocate a single va_range
         // into many range groups.  For this reason, we iterate over each va_range first
         // then through the range groups within.
@@ -637,6 +647,8 @@ static NV_STATUS uvm_migrate(uvm_va_space_t *va_space,
 
     if (mm)
         uvm_assert_mmap_lock_locked(mm);
+    else if (!first_va_range)
+        return NV_ERR_INVALID_ADDRESS;
 
     va_block_context = uvm_va_block_context_alloc(mm);
     if (!va_block_context)

kernel-open/nvidia-uvm/uvm_migrate_pageable.c

@@ -836,6 +836,17 @@ static NV_STATUS migrate_pageable_vma_region(struct vm_area_struct *vma,
     return NV_OK;
 }
 
+NV_STATUS uvm_test_skip_migrate_vma(UVM_TEST_SKIP_MIGRATE_VMA_PARAMS *params, struct file *filp)
+{
+    uvm_va_space_t *va_space = uvm_va_space_get(filp);
+
+    uvm_va_space_down_write(va_space);
+    va_space->test.skip_migrate_vma = params->skip;
+    uvm_va_space_up_write(va_space);
+
+    return NV_OK;
+}
+
 static NV_STATUS migrate_pageable_vma(struct vm_area_struct *vma,
                                       unsigned long start,
                                       unsigned long outer,
@@ -858,6 +869,9 @@ static NV_STATUS migrate_pageable_vma(struct vm_area_struct *vma,
     start = max(start, vma->vm_start);
     outer = min(outer, vma->vm_end);
 
+    if (va_space->test.skip_migrate_vma)
+        return NV_WARN_NOTHING_TO_DO;
+
     // TODO: Bug 2419180: support file-backed pages in migrate_vma, when
     //       support for it is added to the Linux kernel
     if (!vma_is_anonymous(vma))
@@ -920,7 +934,9 @@ static NV_STATUS migrate_pageable(migrate_vma_state_t *state)
             bool touch = uvm_migrate_args->touch;
             uvm_populate_permissions_t populate_permissions = uvm_migrate_args->populate_permissions;
 
-            UVM_ASSERT(!vma_is_anonymous(vma) || uvm_processor_mask_empty(&va_space->registered_gpus));
+            UVM_ASSERT(va_space->test.skip_migrate_vma ||
+                       !vma_is_anonymous(vma) ||
+                       uvm_processor_mask_empty(&va_space->registered_gpus));
 
             // We can't use migrate_vma to move the pages as desired. Normally
             // this fallback path is supposed to populate the memory then inform

kernel-open/nvidia-uvm/uvm_migrate_pageable.h

@@ -34,8 +34,8 @@ typedef struct
 {
     uvm_va_space_t                  *va_space;
     struct mm_struct                *mm;
-    unsigned long                   start;
-    unsigned long                   length;
+    const unsigned long             start;
+    const unsigned long             length;
     uvm_processor_id_t              dst_id;
 
     // dst_node_id may be clobbered by uvm_migrate_pageable().
@@ -51,7 +51,7 @@ typedef struct
 #if defined(CONFIG_MIGRATE_VMA_HELPER)
 #define UVM_MIGRATE_VMA_SUPPORTED 1
 #else
-#if defined(CONFIG_DEVICE_PRIVATE) && defined(NV_MIGRATE_VMA_SETUP_PRESENT)
+#if NV_IS_EXPORT_SYMBOL_PRESENT_migrate_vma_setup
 #define UVM_MIGRATE_VMA_SUPPORTED 1
 #endif
 #endif
@@ -218,6 +218,9 @@ NV_STATUS uvm_migrate_pageable(uvm_migrate_args_t *uvm_migrate_args);
 NV_STATUS uvm_migrate_pageable_init(void);
 
 void uvm_migrate_pageable_exit(void);
+
+NV_STATUS uvm_test_skip_migrate_vma(UVM_TEST_SKIP_MIGRATE_VMA_PARAMS *params, struct file *filp);
+
 #else // UVM_MIGRATE_VMA_SUPPORTED
 
 static NV_STATUS uvm_migrate_pageable(uvm_migrate_args_t *uvm_migrate_args)
@@ -251,6 +254,10 @@ static void uvm_migrate_pageable_exit(void)
 {
 }
 
+static inline NV_STATUS uvm_test_skip_migrate_vma(UVM_TEST_SKIP_MIGRATE_VMA_PARAMS *params, struct file *filp)
+{
+    return NV_OK;
+}
 #endif // UVM_MIGRATE_VMA_SUPPORTED
 
 #endif

kernel-open/nvidia-uvm/uvm_mmu.c

@@ -323,37 +323,156 @@ static void uvm_mmu_page_table_cpu_memset_16(uvm_gpu_t *gpu,
     uvm_mmu_page_table_cpu_unmap(gpu, phys_alloc);
 }
 
+static void pde_fill_cpu(uvm_page_tree_t *tree,
+                         uvm_page_directory_t *directory,
+                         NvU32 start_index,
+                         NvU32 pde_count,
+                         uvm_mmu_page_table_alloc_t **phys_addr)
+{
+    NvU64 pde_data[2], entry_size;
+    NvU32 i;
+
+    UVM_ASSERT(uvm_mmu_use_cpu(tree));
+
+    entry_size = tree->hal->entry_size(directory->depth);
+    UVM_ASSERT(sizeof(pde_data) >= entry_size);
+
+    for (i = 0; i < pde_count; i++) {
+        tree->hal->make_pde(pde_data, phys_addr, directory, start_index + i);
+
+        if (entry_size == sizeof(pde_data[0]))
+            uvm_mmu_page_table_cpu_memset_8(tree->gpu, &directory->phys_alloc, start_index + i, pde_data[0], 1);
+        else
+            uvm_mmu_page_table_cpu_memset_16(tree->gpu, &directory->phys_alloc, start_index + i, pde_data, 1);
+    }
+}
+
+static void pde_fill_gpu(uvm_page_tree_t *tree,
+                         uvm_page_directory_t *directory,
+                         NvU32 start_index,
+                         NvU32 pde_count,
+                         uvm_mmu_page_table_alloc_t **phys_addr,
+                         uvm_push_t *push)
+{
+    NvU64 pde_data[2], entry_size;
+    uvm_gpu_address_t pde_entry_addr = uvm_mmu_gpu_address(tree->gpu, directory->phys_alloc.addr);
+    NvU32 max_inline_entries;
+    uvm_push_flag_t push_membar_flag = UVM_PUSH_FLAG_COUNT;
+    uvm_gpu_address_t inline_data_addr;
+    uvm_push_inline_data_t inline_data;
+    NvU32 entry_count, i, j;
+
+    UVM_ASSERT(!uvm_mmu_use_cpu(tree));
+
+    entry_size = tree->hal->entry_size(directory->depth);
+    UVM_ASSERT(sizeof(pde_data) >= entry_size);
+
+    max_inline_entries = UVM_PUSH_INLINE_DATA_MAX_SIZE / entry_size;
+
+    if (uvm_push_get_and_reset_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE))
+        push_membar_flag = UVM_PUSH_FLAG_NEXT_MEMBAR_NONE;
+    else if (uvm_push_get_and_reset_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_GPU))
+        push_membar_flag = UVM_PUSH_FLAG_NEXT_MEMBAR_GPU;
+
+    pde_entry_addr.address += start_index * entry_size;
+
+    for (i = 0; i < pde_count;) {
+        // All but the first memory operation can be pipelined. We respect the
+        // caller's pipelining settings for the first push.
+        if (i != 0)
+            uvm_push_set_flag(push, UVM_PUSH_FLAG_CE_NEXT_PIPELINED);
+
+        entry_count = min(pde_count - i, max_inline_entries);
+
+        // No membar is needed until the last memory operation. Otherwise,
+        // use caller's membar flag.
+        if ((i + entry_count) < pde_count)
+            uvm_push_set_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE);
+        else if (push_membar_flag != UVM_PUSH_FLAG_COUNT)
+            uvm_push_set_flag(push, push_membar_flag);
+
+        uvm_push_inline_data_begin(push, &inline_data);
+        for (j = 0; j < entry_count; j++) {
+            tree->hal->make_pde(pde_data, phys_addr, directory, start_index + i + j);
+            uvm_push_inline_data_add(&inline_data, pde_data, entry_size);
+        }
+        inline_data_addr = uvm_push_inline_data_end(&inline_data);
+
+        tree->gpu->parent->ce_hal->memcopy(push, pde_entry_addr, inline_data_addr, entry_count * entry_size);
+
+        i += entry_count;
+        pde_entry_addr.address += entry_size * entry_count;
+    }
+}
+
+// pde_fill() populates pde_count PDE entries (starting at start_index) with
+// the same mapping, i.e., with the same physical address (phys_addr).
+// pde_fill() is optimized for pde_count == 1, which is the common case.
+static void pde_fill(uvm_page_tree_t *tree,
+                     uvm_page_directory_t *directory,
+                     NvU32 start_index,
+                     NvU32 pde_count,
+                     uvm_mmu_page_table_alloc_t **phys_addr,
+                     uvm_push_t *push)
+{
+    UVM_ASSERT(start_index + pde_count <= uvm_mmu_page_tree_entries(tree, directory->depth, UVM_PAGE_SIZE_AGNOSTIC));
+
+    if (push)
+        pde_fill_gpu(tree, directory, start_index, pde_count, phys_addr, push);
+    else
+        pde_fill_cpu(tree, directory, start_index, pde_count, phys_addr);
+}
+
 static void phys_mem_init(uvm_page_tree_t *tree, NvU32 page_size, uvm_page_directory_t *dir, uvm_push_t *push)
 {
-    NvU64 clear_bits[2];
-    uvm_mmu_mode_hal_t *hal = tree->hal;
+    NvU32 entries_count = uvm_mmu_page_tree_entries(tree, dir->depth, page_size);
+    NvU8 max_pde_depth = tree->hal->page_table_depth(UVM_PAGE_SIZE_AGNOSTIC) - 1;
 
-    if (dir->depth == tree->hal->page_table_depth(page_size)) {
-        *clear_bits = 0; // Invalid PTE
-    }
-    else {
-        // passing in NULL for the phys_allocs will mark the child entries as invalid
-        uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
-        hal->make_pde(clear_bits, phys_allocs, dir->depth);
+    // Passing in NULL for the phys_allocs will mark the child entries as
+    // invalid.
+    uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
 
-        // Make sure that using only clear_bits[0] will work
-        UVM_ASSERT(hal->entry_size(dir->depth) == sizeof(clear_bits[0]) || clear_bits[0] == clear_bits[1]);
-    }
+    // Init with an invalid PTE or clean PDE. Only Maxwell PDEs can have more
+    // than 512 entries. In this case, we initialize them all with the same
+    // clean PDE. ATS systems may require clean PDEs with
+    // ATS_ALLOWED/ATS_NOT_ALLOWED bit settings based on the mapping VA.
+    // We only clean_bits to 0 at the lowest page table level (PTE table), i.e.,
+    // when depth is greater than the max_pde_depth.
+    if ((dir->depth > max_pde_depth) || (entries_count > 512 && !g_uvm_global.ats.enabled)) {
+        NvU64 clear_bits[2];
 
-    // initialize the memory to a reasonable value
-    if (push) {
-        tree->gpu->parent->ce_hal->memset_8(push,
-                                            uvm_mmu_gpu_address(tree->gpu, dir->phys_alloc.addr),
+        // If it is not a PTE, make a clean PDE.
+        if (dir->depth != tree->hal->page_table_depth(page_size)) {
+            // make_pde() child index is zero/ignored, since it is only used in
+            // PDEs on ATS-enabled systems where pde_fill() is preferred.
+            tree->hal->make_pde(clear_bits, phys_allocs, dir, 0);
+
+            // Make sure that using only clear_bits[0] will work.
+            UVM_ASSERT(tree->hal->entry_size(dir->depth) == sizeof(clear_bits[0]) || clear_bits[0] == clear_bits[1]);
+        }
+        else {
+            *clear_bits = 0;
+        }
+
+        // Initialize the memory to a reasonable value.
+        if (push) {
+            tree->gpu->parent->ce_hal->memset_8(push,
+                                                uvm_mmu_gpu_address(tree->gpu, dir->phys_alloc.addr),
+                                                *clear_bits,
+                                                dir->phys_alloc.size);
+        }
+        else {
+            uvm_mmu_page_table_cpu_memset_8(tree->gpu,
+                                            &dir->phys_alloc,
+                                            0,
                                             *clear_bits,
-                                            dir->phys_alloc.size);
+                                            dir->phys_alloc.size / sizeof(*clear_bits));
+        }
     }
     else {
-        uvm_mmu_page_table_cpu_memset_8(tree->gpu,
-                                        &dir->phys_alloc,
-                                        0,
-                                        *clear_bits,
-                                        dir->phys_alloc.size / sizeof(*clear_bits));
+        pde_fill(tree, dir, 0, entries_count, phys_allocs, push);
     }
+
 }
 
 static uvm_page_directory_t *allocate_directory(uvm_page_tree_t *tree,
@@ -367,8 +486,10 @@ static uvm_page_directory_t *allocate_directory(uvm_page_tree_t *tree,
     NvLength phys_alloc_size = hal->allocation_size(depth, page_size);
     uvm_page_directory_t *dir;
 
-    // The page tree doesn't cache PTEs so space is not allocated for entries that are always PTEs.
-    // 2M PTEs may later become PDEs so pass UVM_PAGE_SIZE_AGNOSTIC, not page_size.
+    // The page tree doesn't cache PTEs so space is not allocated for entries
+    // that are always PTEs.
+    // 2M PTEs may later become PDEs so pass UVM_PAGE_SIZE_AGNOSTIC, not
+    // page_size.
     if (depth == hal->page_table_depth(UVM_PAGE_SIZE_AGNOSTIC))
         entry_count = 0;
     else
@@ -409,108 +530,6 @@ static inline NvU32 index_to_entry(uvm_mmu_mode_hal_t *hal, NvU32 entry_index, N
     return hal->entries_per_index(depth) * entry_index + hal->entry_offset(depth, page_size);
 }
 
-static void pde_fill_cpu(uvm_page_tree_t *tree,
-                         NvU32 depth,
-                         uvm_mmu_page_table_alloc_t *directory,
-                         NvU32 start_index,
-                         NvU32 pde_count,
-                         uvm_mmu_page_table_alloc_t **phys_addr)
-{
-    NvU64 pde_data[2], entry_size;
-
-    UVM_ASSERT(uvm_mmu_use_cpu(tree));
-    entry_size = tree->hal->entry_size(depth);
-    UVM_ASSERT(sizeof(pde_data) >= entry_size);
-
-    tree->hal->make_pde(pde_data, phys_addr, depth);
-
-    if (entry_size == sizeof(pde_data[0]))
-        uvm_mmu_page_table_cpu_memset_8(tree->gpu, directory, start_index, pde_data[0], pde_count);
-    else
-        uvm_mmu_page_table_cpu_memset_16(tree->gpu, directory, start_index, pde_data, pde_count);
-}
-
-static void pde_fill_gpu(uvm_page_tree_t *tree,
-                         NvU32 depth,
-                         uvm_mmu_page_table_alloc_t *directory,
-                         NvU32 start_index,
-                         NvU32 pde_count,
-                         uvm_mmu_page_table_alloc_t **phys_addr,
-                         uvm_push_t *push)
-{
-    NvU64 pde_data[2], entry_size;
-    uvm_gpu_address_t pde_entry_addr = uvm_mmu_gpu_address(tree->gpu, directory->addr);
-
-    UVM_ASSERT(!uvm_mmu_use_cpu(tree));
-
-    entry_size = tree->hal->entry_size(depth);
-    UVM_ASSERT(sizeof(pde_data) >= entry_size);
-
-    tree->hal->make_pde(pde_data, phys_addr, depth);
-    pde_entry_addr.address += start_index * entry_size;
-
-    if (entry_size == sizeof(pde_data[0])) {
-        tree->gpu->parent->ce_hal->memset_8(push, pde_entry_addr, pde_data[0], sizeof(pde_data[0]) * pde_count);
-    }
-    else {
-        NvU32 max_inline_entries = UVM_PUSH_INLINE_DATA_MAX_SIZE / sizeof(pde_data);
-        uvm_gpu_address_t inline_data_addr;
-        uvm_push_inline_data_t inline_data;
-        NvU32 membar_flag = 0;
-        NvU32 i;
-
-        if (uvm_push_get_and_reset_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE))
-            membar_flag = UVM_PUSH_FLAG_NEXT_MEMBAR_NONE;
-        else if (uvm_push_get_and_reset_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_GPU))
-            membar_flag = UVM_PUSH_FLAG_NEXT_MEMBAR_GPU;
-
-        for (i = 0; i < pde_count;) {
-            NvU32 j;
-            NvU32 entry_count = min(pde_count - i, max_inline_entries);
-
-            uvm_push_inline_data_begin(push, &inline_data);
-            for (j = 0; j < entry_count; j++)
-                uvm_push_inline_data_add(&inline_data, pde_data, sizeof(pde_data));
-            inline_data_addr = uvm_push_inline_data_end(&inline_data);
-
-            // All but the first memcopy can be pipelined. We respect the
-            // caller's pipelining settings for the first push.
-            if (i != 0)
-                uvm_push_set_flag(push, UVM_PUSH_FLAG_CE_NEXT_PIPELINED);
-
-            // No membar is needed until the last copy. Otherwise, use
-            // caller's membar flag.
-            if (i + entry_count < pde_count)
-                uvm_push_set_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE);
-            else if (membar_flag)
-                uvm_push_set_flag(push, membar_flag);
-
-            tree->gpu->parent->ce_hal->memcopy(push, pde_entry_addr, inline_data_addr, entry_count * sizeof(pde_data));
-
-            i += entry_count;
-            pde_entry_addr.address += sizeof(pde_data) * entry_count;
-        }
-    }
-}
-
-// pde_fill() populates pde_count PDE entries (starting at start_index) with
-// the same mapping, i.e., with the same physical address (phys_addr).
-static void pde_fill(uvm_page_tree_t *tree,
-                     NvU32 depth,
-                     uvm_mmu_page_table_alloc_t *directory,
-                     NvU32 start_index,
-                     NvU32 pde_count,
-                     uvm_mmu_page_table_alloc_t **phys_addr,
-                     uvm_push_t *push)
-{
-    UVM_ASSERT(start_index + pde_count <= uvm_mmu_page_tree_entries(tree, depth, UVM_PAGE_SIZE_AGNOSTIC));
-
-    if (push)
-        pde_fill_gpu(tree, depth, directory, start_index, pde_count, phys_addr, push);
-    else
-        pde_fill_cpu(tree, depth, directory, start_index, pde_count, phys_addr);
-}
-
 static uvm_page_directory_t *host_pde_write(uvm_page_directory_t *dir,
                                             uvm_page_directory_t *parent,
                                             NvU32 index_in_parent)
@@ -540,7 +559,7 @@ static void pde_write(uvm_page_tree_t *tree,
             phys_allocs[i] = &entry->phys_alloc;
     }
 
-    pde_fill(tree, dir->depth, &dir->phys_alloc, entry_index, 1, phys_allocs, push);
+    pde_fill(tree, dir, entry_index, 1, phys_allocs, push);
 }
 
 static void host_pde_clear(uvm_page_tree_t *tree, uvm_page_directory_t *dir, NvU32 entry_index, NvU32 page_size)
@@ -800,7 +819,6 @@ static void free_unused_directories(uvm_page_tree_t *tree,
             }
         }
     }
-
 }
 
 static NV_STATUS allocate_page_table(uvm_page_tree_t *tree, NvU32 page_size, uvm_mmu_page_table_alloc_t *out)
@@ -811,10 +829,93 @@ static NV_STATUS allocate_page_table(uvm_page_tree_t *tree, NvU32 page_size, uvm
     return phys_mem_allocate(tree, alloc_size, tree->location, UVM_PMM_ALLOC_FLAGS_EVICT, out);
 }
 
+static bool page_tree_ats_init_required(uvm_page_tree_t *tree)
+{
+    // We have full control of the kernel page tables mappings, no ATS address
+    // aliases is expected.
+    if (tree->type == UVM_PAGE_TREE_TYPE_KERNEL)
+        return false;
+
+    // Enable uvm_page_tree_init() from the page_tree test.
+    if (uvm_enable_builtin_tests && tree->gpu_va_space == NULL)
+        return false;
+
+    if (!tree->gpu_va_space->ats.enabled)
+        return false;
+
+    return tree->gpu->parent->no_ats_range_required;
+}
+
+static NV_STATUS page_tree_ats_init(uvm_page_tree_t *tree)
+{
+    NV_STATUS status;
+    NvU64 min_va_upper, max_va_lower;
+    NvU32 page_size;
+
+    if (!page_tree_ats_init_required(tree))
+        return NV_OK;
+
+    page_size = uvm_mmu_biggest_page_size(tree);
+
+    uvm_cpu_get_unaddressable_range(&max_va_lower, &min_va_upper);
+
+    // Potential violation of the UVM internal get/put_ptes contract. get_ptes()
+    // creates and initializes enough PTEs to populate all PDEs covering the
+    // no_ats_ranges. We store the no_ats_ranges in the tree, so they can be
+    // put_ptes()'ed on deinit(). It doesn't preclude the range to be used by a
+    // future get_ptes(), since we don't write to the PTEs (range->table) from
+    // the tree->no_ats_ranges.
+    //
+    // Lower half
+    status = uvm_page_tree_get_ptes(tree,
+                                    page_size,
+                                    max_va_lower,
+                                    page_size,
+                                    UVM_PMM_ALLOC_FLAGS_EVICT,
+                                    &tree->no_ats_ranges[0]);
+    if (status != NV_OK)
+        return status;
+
+    UVM_ASSERT(tree->no_ats_ranges[0].entry_count == 1);
+
+    if (uvm_platform_uses_canonical_form_address()) {
+        // Upper half
+        status = uvm_page_tree_get_ptes(tree,
+                                        page_size,
+                                        min_va_upper - page_size,
+                                        page_size,
+                                        UVM_PMM_ALLOC_FLAGS_EVICT,
+                                        &tree->no_ats_ranges[1]);
+        if (status != NV_OK)
+            return status;
+
+        UVM_ASSERT(tree->no_ats_ranges[1].entry_count == 1);
+    }
+
+    return NV_OK;
+}
+
+static void page_tree_ats_deinit(uvm_page_tree_t *tree)
+{
+    size_t i;
+
+    if (page_tree_ats_init_required(tree)) {
+        for (i = 0; i < ARRAY_SIZE(tree->no_ats_ranges); i++) {
+            if (tree->no_ats_ranges[i].entry_count)
+                uvm_page_tree_put_ptes(tree, &tree->no_ats_ranges[i]);
+        }
+
+        memset(tree->no_ats_ranges, 0, sizeof(tree->no_ats_ranges));
+    }
+}
+
 static void map_remap_deinit(uvm_page_tree_t *tree)
 {
-    if (tree->map_remap.pde0.size)
-        phys_mem_deallocate(tree, &tree->map_remap.pde0);
+    if (tree->map_remap.pde0) {
+        phys_mem_deallocate(tree, &tree->map_remap.pde0->phys_alloc);
+        uvm_kvfree(tree->map_remap.pde0);
+        tree->map_remap.pde0 = NULL;
+    }
 
     if (tree->map_remap.ptes_invalid_4k.size)
         phys_mem_deallocate(tree, &tree->map_remap.ptes_invalid_4k);
@@ -839,10 +940,16 @@ static NV_STATUS map_remap_init(uvm_page_tree_t *tree)
     // PDE1-depth(512M) PTE. We first map it to the pde0 directory, then we
     // return the PTE for the get_ptes()'s caller.
     if (tree->hal->page_sizes() & UVM_PAGE_SIZE_512M) {
-        status = allocate_page_table(tree, UVM_PAGE_SIZE_2M, &tree->map_remap.pde0);
-        if (status != NV_OK)
+        tree->map_remap.pde0 = allocate_directory(tree,
+                                                  UVM_PAGE_SIZE_2M,
+                                                  tree->hal->page_table_depth(UVM_PAGE_SIZE_2M),
+                                                  UVM_PMM_ALLOC_FLAGS_EVICT);
+        if (tree->map_remap.pde0 == NULL) {
+            status = NV_ERR_NO_MEMORY;
             goto error;
+        }
     }
+
     status = page_tree_begin_acquire(tree, &tree->tracker, &push, "map remap init");
     if (status != NV_OK)
         goto error;
@@ -864,22 +971,23 @@ static NV_STATUS map_remap_init(uvm_page_tree_t *tree)
         uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
         NvU32 depth = tree->hal->page_table_depth(UVM_PAGE_SIZE_4K) - 1;
         size_t index_4k = tree->hal->entry_offset(depth, UVM_PAGE_SIZE_4K);
-
-        // pde0 depth equals UVM_PAGE_SIZE_2M.
-        NvU32 pde0_depth = tree->hal->page_table_depth(UVM_PAGE_SIZE_2M);
-        NvU32 pde0_entries = tree->map_remap.pde0.size / tree->hal->entry_size(pde0_depth);
+        NvU32 pde0_entries = tree->map_remap.pde0->phys_alloc.size / tree->hal->entry_size(tree->map_remap.pde0->depth);
 
         // The big-page entry is NULL which makes it an invalid entry.
         phys_allocs[index_4k] = &tree->map_remap.ptes_invalid_4k;
 
         // By default CE operations include a MEMBAR_SYS. MEMBAR_GPU is
         // sufficient when pde0 is allocated in VIDMEM.
-        if (tree->map_remap.pde0.addr.aperture == UVM_APERTURE_VID)
+        if (tree->map_remap.pde0->phys_alloc.addr.aperture == UVM_APERTURE_VID)
             uvm_push_set_flag(&push, UVM_PUSH_FLAG_NEXT_MEMBAR_GPU);
 
+        // This is an orphan directory, make_pde() requires a directory to
+        // compute the VA. The UVM depth map_remap() operates on is not in the
+        // range make_pde() must operate. We only need to supply the fields used
+        // by make_pde() to not access invalid memory addresses.
+
         pde_fill(tree,
-                 pde0_depth,
-                 &tree->map_remap.pde0,
+                 tree->map_remap.pde0,
                  0,
                  pde0_entries,
                  (uvm_mmu_page_table_alloc_t **)&phys_allocs,
@@ -906,11 +1014,10 @@ error:
 // --------------|-------------------------||----------------|----------------
 //    vidmem     |           -             ||    vidmem      |      false
 //    sysmem     |           -             ||    sysmem      |      false
-//    default    |        <not set>        ||    vidmem      |      true (1)
+//    default    |        <not set>        ||    vidmem      |      true
 //    default    |         vidmem          ||    vidmem      |      false
 //    default    |         sysmem          ||    sysmem      |      false
 //
-// (1) When SEV mode is enabled, the fallback path is disabled.
 //
 // In SR-IOV heavy the the page tree must be in vidmem, to prevent guest drivers
 // from updating GPU page tables without hypervisor knowledge.
@@ -926,28 +1033,27 @@ error:
 //
 static void page_tree_set_location(uvm_page_tree_t *tree, uvm_aperture_t location)
 {
-    bool should_location_be_vidmem;
     UVM_ASSERT(tree->gpu != NULL);
     UVM_ASSERT_MSG((location == UVM_APERTURE_VID) ||
                    (location == UVM_APERTURE_SYS) ||
                    (location == UVM_APERTURE_DEFAULT),
                    "Invalid location %s (%d)\n", uvm_aperture_string(location), (int)location);
 
-    should_location_be_vidmem = uvm_gpu_is_virt_mode_sriov_heavy(tree->gpu)
-                                || uvm_conf_computing_mode_enabled(tree->gpu);
-
     // The page tree of a "fake" GPU used during page tree testing can be in
-    // sysmem even if should_location_be_vidmem is true. A fake GPU can be
-    // identified by having no channel manager.
-    if ((tree->gpu->channel_manager != NULL) && should_location_be_vidmem)
-        UVM_ASSERT(location == UVM_APERTURE_VID);
+    // sysmem in scenarios where a "real" GPU must be in vidmem. Fake GPUs can
+    // be identified by having no channel manager.
+    if (tree->gpu->channel_manager != NULL) {
+
+        if (uvm_gpu_is_virt_mode_sriov_heavy(tree->gpu))
+            UVM_ASSERT(location == UVM_APERTURE_VID);
+        else if (uvm_conf_computing_mode_enabled(tree->gpu))
+            UVM_ASSERT(location == UVM_APERTURE_VID);
+    }
 
     if (location == UVM_APERTURE_DEFAULT) {
         if (page_table_aperture == UVM_APERTURE_DEFAULT) {
             tree->location = UVM_APERTURE_VID;
-
-            // See the comment (1) above.
-            tree->location_sys_fallback = !g_uvm_global.sev_enabled;
+            tree->location_sys_fallback = true;
         }
         else {
             tree->location = page_table_aperture;
@@ -1008,11 +1114,22 @@ NV_STATUS uvm_page_tree_init(uvm_gpu_t *gpu,
         return status;
 
     phys_mem_init(tree, UVM_PAGE_SIZE_AGNOSTIC, tree->root, &push);
-    return page_tree_end_and_wait(tree, &push);
+
+    status = page_tree_end_and_wait(tree, &push);
+    if (status != NV_OK)
+        return status;
+
+    status = page_tree_ats_init(tree);
+    if (status != NV_OK)
+        return status;
+
+    return NV_OK;
 }
 
 void uvm_page_tree_deinit(uvm_page_tree_t *tree)
 {
+    page_tree_ats_deinit(tree);
+
     UVM_ASSERT(tree->root->ref_count == 0);
 
     // Take the tree lock only to avoid assertions. It is not required for
@@ -1251,7 +1368,6 @@ static NV_STATUS try_get_ptes(uvm_page_tree_t *tree,
         UVM_ASSERT(uvm_gpu_can_address_kernel(tree->gpu, start, size));
 
     while (true) {
-
         // index of the entry, for the first byte of the range, within its
         // containing directory
         NvU32 start_index;
@@ -1283,7 +1399,8 @@ static NV_STATUS try_get_ptes(uvm_page_tree_t *tree,
                 if (dir_cache[dir->depth] == NULL) {
                     *cur_depth = dir->depth;
 
-                    // Undo the changes to the tree so that the dir cache remains private to the thread
+                    // Undo the changes to the tree so that the dir cache
+                    // remains private to the thread.
                     for (i = 0; i < used_count; i++)
                         host_pde_clear(tree, dirs_used[i]->host_parent, dirs_used[i]->index_in_parent, page_size);
 
@@ -1334,10 +1451,9 @@ static NV_STATUS map_remap(uvm_page_tree_t *tree, NvU64 start, NvLength size, uv
     if (uvm_page_table_range_aperture(range) == UVM_APERTURE_VID)
         uvm_push_set_flag(&push, UVM_PUSH_FLAG_NEXT_MEMBAR_GPU);
 
-    phys_alloc[0] = &tree->map_remap.pde0;
+    phys_alloc[0] = &tree->map_remap.pde0->phys_alloc;
     pde_fill(tree,
-             range->table->depth,
-             &range->table->phys_alloc,
+             range->table,
              range->start_index,
              range->entry_count,
              (uvm_mmu_page_table_alloc_t **)&phys_alloc,
@@ -1382,7 +1498,8 @@ NV_STATUS uvm_page_tree_get_ptes_async(uvm_page_tree_t *tree,
                                   dir_cache)) == NV_ERR_MORE_PROCESSING_REQUIRED) {
         uvm_mutex_unlock(&tree->lock);
 
-        // try_get_ptes never needs depth 0, so store a directory at its parent's depth
+        // try_get_ptes never needs depth 0, so store a directory at its
+        // parent's depth.
         // TODO: Bug 1766655: Allocate everything below cur_depth instead of
         //       retrying for every level.
         dir_cache[cur_depth] = allocate_directory(tree, page_size, cur_depth + 1, pmm_flags);
@@ -1665,8 +1782,12 @@ NV_STATUS uvm_page_table_range_vec_init(uvm_page_tree_t *tree,
                                               range);
         if (status != NV_OK) {
             UVM_ERR_PRINT("Failed to get PTEs for subrange %zd [0x%llx, 0x%llx) size 0x%llx, part of [0x%llx, 0x%llx)\n",
-                    i, range_start, range_start + range_size, range_size,
-                    start, size);
+                          i,
+                          range_start,
+                          range_start + range_size,
+                          range_size,
+                          start,
+                          size);
             goto out;
         }
     }

kernel-open/nvidia-uvm/uvm_mmu.h

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2022 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -162,7 +162,7 @@ struct uvm_page_directory_struct
     // pointers to child directories on the host.
     // this array is variable length, so it needs to be last to allow it to
     // take up extra space
-    uvm_page_directory_t *entries[0];
+    uvm_page_directory_t *entries[];
 };
 
 enum
@@ -215,11 +215,14 @@ struct uvm_mmu_mode_hal_struct
     // memory out-of-range error so we can immediately identify bad PTE usage.
     NvU64 (*poisoned_pte)(void);
 
-    // write a PDE bit-pattern to entry based on the data in entries (which may
+    // Write a PDE bit-pattern to entry based on the data in allocs (which may
     // point to two items for dual PDEs).
-    // any of allocs are allowed to be NULL, in which case they are to be
-    // treated as empty.
-    void (*make_pde)(void *entry, uvm_mmu_page_table_alloc_t **allocs, NvU32 depth);
+    // Any of allocs are allowed to be NULL, in which case they are to be
+    // treated as empty. make_pde() uses dir and child_index to compute the
+    // mapping PDE VA. On ATS-enabled systems, we may set PDE's PCF as
+    // ATS_ALLOWED or ATS_NOT_ALLOWED based on the mapping PDE VA, even for
+    // invalid/clean PDE entries.
+    void (*make_pde)(void *entry, uvm_mmu_page_table_alloc_t **allocs, uvm_page_directory_t *dir, NvU32 child_index);
 
     // size of an entry in a directory/table.  Generally either 8 or 16 bytes.
     // (in the case of Pascal dual PDEs)
@@ -229,7 +232,7 @@ struct uvm_mmu_mode_hal_struct
     NvU32 (*entries_per_index)(NvU32 depth);
 
     // For dual PDEs, this is ether 1 or 0, depending on the page size.
-    // This is used to index the host copy only.  GPU PDEs are always entirely
+    // This is used to index the host copy only. GPU PDEs are always entirely
     // re-written using make_pde.
     NvLength (*entry_offset)(NvU32 depth, NvU32 page_size);
 
@@ -295,11 +298,16 @@ struct uvm_page_tree_struct
 
         // PDE0 where all big-page entries are invalid, and small-page entries
         // point to ptes_invalid_4k.
-        // pde0 is only used on Pascal-Ampere, i.e., they have the same PDE
-        // format.
-        uvm_mmu_page_table_alloc_t pde0;
+        // pde0 is used on Pascal+ GPUs, i.e., they have the same PDE format.
+        uvm_page_directory_t *pde0;
     } map_remap;
 
+    // On ATS-enabled systems where the CPU VA width is smaller than the GPU VA
+    // width, the excess address range is set with ATS_NOT_ALLOWED on all  leaf
+    // PDEs covering that range. We have at most 2 no_ats_ranges, due to
+    // canonical form address systems.
+    uvm_page_table_range_t no_ats_ranges[2];
+
     // Tracker for all GPU operations on the tree
     uvm_tracker_t tracker;
 };
@@ -365,21 +373,32 @@ void uvm_page_tree_deinit(uvm_page_tree_t *tree);
 // the same page size without an intervening put_ptes. To duplicate a subset of
 // an existing range or change the size of an existing range, use
 // uvm_page_table_range_get_upper() and/or uvm_page_table_range_shrink().
-NV_STATUS uvm_page_tree_get_ptes(uvm_page_tree_t *tree, NvU32 page_size, NvU64 start, NvLength size,
-        uvm_pmm_alloc_flags_t pmm_flags, uvm_page_table_range_t *range);
+NV_STATUS uvm_page_tree_get_ptes(uvm_page_tree_t *tree,
+                                 NvU32 page_size,
+                                 NvU64 start,
+                                 NvLength size,
+                                 uvm_pmm_alloc_flags_t pmm_flags,
+                                 uvm_page_table_range_t *range);
 
 // Same as uvm_page_tree_get_ptes(), but doesn't synchronize the GPU work.
 //
 // All pending operations can be waited on with uvm_page_tree_wait().
-NV_STATUS uvm_page_tree_get_ptes_async(uvm_page_tree_t *tree, NvU32 page_size, NvU64 start, NvLength size,
-        uvm_pmm_alloc_flags_t pmm_flags, uvm_page_table_range_t *range);
+NV_STATUS uvm_page_tree_get_ptes_async(uvm_page_tree_t *tree,
+                                       NvU32 page_size,
+                                       NvU64 start,
+                                       NvLength size,
+                                       uvm_pmm_alloc_flags_t pmm_flags,
+                                       uvm_page_table_range_t *range);
 
 // Returns a single-entry page table range for the addresses passed.
 // The size parameter must be a page size supported by this tree.
 // This is equivalent to calling uvm_page_tree_get_ptes() with size equal to
 // page_size.
-NV_STATUS uvm_page_tree_get_entry(uvm_page_tree_t *tree, NvU32 page_size, NvU64 start,
-        uvm_pmm_alloc_flags_t pmm_flags, uvm_page_table_range_t *single);
+NV_STATUS uvm_page_tree_get_entry(uvm_page_tree_t *tree,
+                                  NvU32 page_size,
+                                  NvU64 start,
+                                  uvm_pmm_alloc_flags_t pmm_flags,
+                                  uvm_page_table_range_t *single);
 
 // For a single-entry page table range, write the PDE (which could be a dual
 // PDE) to the GPU.
@@ -478,8 +497,8 @@ NV_STATUS uvm_page_table_range_vec_create(uvm_page_tree_t *tree,
 // new_range_vec will contain the upper portion of range_vec, starting at
 // new_end + 1.
 //
-// new_end + 1 is required to be within the address range of range_vec and be aligned to
-// range_vec's page_size.
+// new_end + 1 is required to be within the address range of range_vec and be
+// aligned to range_vec's page_size.
 //
 // On failure, the original range vector is left unmodified.
 NV_STATUS uvm_page_table_range_vec_split_upper(uvm_page_table_range_vec_t *range_vec,
@@ -501,18 +520,22 @@ void uvm_page_table_range_vec_destroy(uvm_page_table_range_vec_t *range_vec);
 // for each offset.
 // The caller_data pointer is what the caller passed in as caller_data to
 // uvm_page_table_range_vec_write_ptes().
-typedef NvU64 (*uvm_page_table_range_pte_maker_t)(uvm_page_table_range_vec_t *range_vec, NvU64 offset,
-        void *caller_data);
+typedef NvU64 (*uvm_page_table_range_pte_maker_t)(uvm_page_table_range_vec_t *range_vec,
+                                                  NvU64 offset,
+                                                  void *caller_data);
 
-// Write all PTEs covered by the range vector using the given PTE making function.
+// Write all PTEs covered by the range vector using the given PTE making
+// function.
 //
 // After writing all the PTEs a TLB invalidate operation is performed including
 // the passed in tlb_membar.
 //
 // See comments about uvm_page_table_range_pte_maker_t for details about the
 // PTE making callback.
-NV_STATUS uvm_page_table_range_vec_write_ptes(uvm_page_table_range_vec_t *range_vec, uvm_membar_t tlb_membar,
-        uvm_page_table_range_pte_maker_t pte_maker, void *caller_data);
+NV_STATUS uvm_page_table_range_vec_write_ptes(uvm_page_table_range_vec_t *range_vec,
+                                              uvm_membar_t tlb_membar,
+                                              uvm_page_table_range_pte_maker_t pte_maker,
+                                              void *caller_data);
 
 // Set all PTEs covered by the range vector to an empty PTE
 //
@@ -636,8 +659,9 @@ static NvU64 uvm_page_table_range_size(uvm_page_table_range_t *range)
 
 // Get the physical address of the entry at entry_index within the range
 // (counted from range->start_index).
-static uvm_gpu_phys_address_t uvm_page_table_range_entry_address(uvm_page_tree_t *tree, uvm_page_table_range_t *range,
-        size_t entry_index)
+static uvm_gpu_phys_address_t uvm_page_table_range_entry_address(uvm_page_tree_t *tree,
+                                                                 uvm_page_table_range_t *range,
+                                                                 size_t entry_index)
 {
     NvU32 entry_size = uvm_mmu_pte_size(tree, range->page_size);
     uvm_gpu_phys_address_t entry = range->table->phys_alloc.addr;

kernel-open/nvidia-uvm/uvm_page_tree_test.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2022 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -146,9 +146,15 @@ static void fake_tlb_invals_disable(void)
     g_fake_tlb_invals_tracking_enabled = false;
 }
 
-// Fake TLB invalidate VA that just saves off the parameters so that they can be verified later
-static void fake_tlb_invalidate_va(uvm_push_t *push, uvm_gpu_phys_address_t pdb,
-        NvU32 depth, NvU64 base, NvU64 size, NvU32 page_size, uvm_membar_t membar)
+// Fake TLB invalidate VA that just saves off the parameters so that they can be
+// verified later.
+static void fake_tlb_invalidate_va(uvm_push_t *push,
+                                   uvm_gpu_phys_address_t pdb,
+                                   NvU32 depth,
+                                   NvU64 base,
+                                   NvU64 size,
+                                   NvU32 page_size,
+                                   uvm_membar_t membar)
 {
     if (!g_fake_tlb_invals_tracking_enabled)
         return;
@@ -210,8 +216,8 @@ static bool assert_and_reset_last_invalidate(NvU32 expected_depth, bool expected
     }
     if ((g_last_fake_inval->membar == UVM_MEMBAR_NONE) == expected_membar) {
         UVM_TEST_PRINT("Expected %s membar, got %s instead\n",
-                expected_membar ? "a" : "no",
-                uvm_membar_string(g_last_fake_inval->membar));
+                       expected_membar ? "a" : "no",
+                       uvm_membar_string(g_last_fake_inval->membar));
         result = false;
     }
 
@@ -230,7 +236,8 @@ static bool assert_last_invalidate_all(NvU32 expected_depth, bool expected_memba
     }
     if (g_last_fake_inval->base != 0 || g_last_fake_inval->size != -1) {
         UVM_TEST_PRINT("Expected invalidate all but got range [0x%llx, 0x%llx) instead\n",
-                g_last_fake_inval->base, g_last_fake_inval->base + g_last_fake_inval->size);
+                       g_last_fake_inval->base,
+                       g_last_fake_inval->base + g_last_fake_inval->size);
         return false;
     }
     if (g_last_fake_inval->depth != expected_depth) {
@@ -247,15 +254,16 @@ static bool assert_invalidate_range_specific(fake_tlb_invalidate_t *inval,
     UVM_ASSERT(g_fake_tlb_invals_tracking_enabled);
 
     if (g_fake_invals_count == 0) {
-        UVM_TEST_PRINT("Expected an invalidate for range [0x%llx, 0x%llx), but got none\n",
-                base, base + size);
+        UVM_TEST_PRINT("Expected an invalidate for range [0x%llx, 0x%llx), but got none\n", base, base + size);
         return false;
     }
 
     if ((inval->base != base || inval->size != size) && inval->base != 0 && inval->size != -1) {
         UVM_TEST_PRINT("Expected invalidate range [0x%llx, 0x%llx), but got range [0x%llx, 0x%llx) instead\n",
-                base, base + size,
-                inval->base, inval->base + inval->size);
+                        base,
+                        base + size,
+                        inval->base,
+                        inval->base + inval->size);
         return false;
     }
     if (inval->depth != expected_depth) {
@@ -270,7 +278,13 @@ static bool assert_invalidate_range_specific(fake_tlb_invalidate_t *inval,
     return true;
 }
 
-static bool assert_invalidate_range(NvU64 base, NvU64 size, NvU32 page_size, bool allow_inval_all, NvU32 range_depth, NvU32 all_depth, bool expected_membar)
+static bool assert_invalidate_range(NvU64 base,
+                                    NvU64 size,
+                                    NvU32 page_size,
+                                    bool allow_inval_all,
+                                    NvU32 range_depth,
+                                    NvU32 all_depth,
+                                    bool expected_membar)
 {
     NvU32 i;
 
@@ -488,7 +502,6 @@ static NV_STATUS alloc_adjacent_pde_64k_memory(uvm_gpu_t *gpu)
     return NV_OK;
 }
 
-
 static NV_STATUS alloc_nearby_pde_64k_memory(uvm_gpu_t *gpu)
 {
     uvm_page_tree_t tree;
@@ -842,6 +855,7 @@ static NV_STATUS get_two_free_apart(uvm_gpu_t *gpu)
     TEST_CHECK_RET(range2.entry_count == 256);
     TEST_CHECK_RET(range2.table->ref_count == 512);
     TEST_CHECK_RET(range1.table == range2.table);
+
     // 4k page is second entry in a dual PDE
     TEST_CHECK_RET(range1.table == tree.root->entries[0]->entries[0]->entries[0]->entries[1]);
     TEST_CHECK_RET(range1.start_index == 256);
@@ -871,6 +885,7 @@ static NV_STATUS get_overlapping_dual_pdes(uvm_gpu_t *gpu)
     MEM_NV_CHECK_RET(test_page_tree_get_ptes(&tree, UVM_PAGE_SIZE_64K, size, size, &range64k), NV_OK);
     TEST_CHECK_RET(range64k.entry_count == 16);
     TEST_CHECK_RET(range64k.table->ref_count == 16);
+
     // 4k page is second entry in a dual PDE
     TEST_CHECK_RET(range64k.table == tree.root->entries[0]->entries[0]->entries[0]->entries[0]);
     TEST_CHECK_RET(range64k.start_index == 16);
@@ -1030,10 +1045,13 @@ static NV_STATUS test_tlb_invalidates(uvm_gpu_t *gpu)
 
     // Depth 4
     NvU64 extent_pte = UVM_PAGE_SIZE_2M;
+
     // Depth 3
     NvU64 extent_pde0 = extent_pte * (1ull << 8);
+
     // Depth 2
     NvU64 extent_pde1 = extent_pde0 * (1ull << 9);
+
     // Depth 1
     NvU64 extent_pde2 = extent_pde1 * (1ull << 9);
 
@@ -1081,7 +1099,11 @@ static NV_STATUS test_tlb_invalidates(uvm_gpu_t *gpu)
     return status;
 }
 
-static NV_STATUS test_tlb_batch_invalidates_case(uvm_page_tree_t *tree, NvU64 base, NvU64 size, NvU32 min_page_size, NvU32 max_page_size)
+static NV_STATUS test_tlb_batch_invalidates_case(uvm_page_tree_t *tree,
+                                                 NvU64 base,
+                                                 NvU64 size,
+                                                 NvU32 min_page_size,
+                                                 NvU32 max_page_size)
 {
     NV_STATUS status = NV_OK;
     uvm_push_t push;
@@ -1205,7 +1227,11 @@ static bool assert_range_vec_ptes(uvm_page_table_range_vec_t *range_vec, bool ex
             NvU64 expected_pte = expecting_cleared ? 0 : range_vec->size + offset;
             if (*pte != expected_pte) {
                 UVM_TEST_PRINT("PTE is 0x%llx instead of 0x%llx for offset 0x%llx within range [0x%llx, 0x%llx)\n",
-                        *pte, expected_pte, offset, range_vec->start, range_vec->size);
+                               *pte,
+                               expected_pte,
+                               offset,
+                               range_vec->start,
+                               range_vec->size);
                 return false;
             }
             offset += range_vec->page_size;
@@ -1226,7 +1252,11 @@ static NV_STATUS test_range_vec_write_ptes(uvm_page_table_range_vec_t *range_vec
     TEST_CHECK_RET(data.status == NV_OK);
     TEST_CHECK_RET(data.count == range_vec->size / range_vec->page_size);
     TEST_CHECK_RET(assert_invalidate_range_specific(g_last_fake_inval,
-            range_vec->start, range_vec->size, range_vec->page_size, page_table_depth, membar != UVM_MEMBAR_NONE));
+                                                    range_vec->start,
+                                                    range_vec->size,
+                                                    range_vec->page_size,
+                                                    page_table_depth,
+                                                    membar != UVM_MEMBAR_NONE));
     TEST_CHECK_RET(assert_range_vec_ptes(range_vec, false));
 
     fake_tlb_invals_disable();
@@ -1249,7 +1279,11 @@ static NV_STATUS test_range_vec_clear_ptes(uvm_page_table_range_vec_t *range_vec
     return NV_OK;
 }
 
-static NV_STATUS test_range_vec_create(uvm_page_tree_t *tree, NvU64 start, NvU64 size, NvU32 page_size, uvm_page_table_range_vec_t **range_vec_out)
+static NV_STATUS test_range_vec_create(uvm_page_tree_t *tree,
+                                       NvU64 start,
+                                       NvU64 size,
+                                       NvU32 page_size,
+                                       uvm_page_table_range_vec_t **range_vec_out)
 {
     uvm_page_table_range_vec_t *range_vec;
     uvm_pmm_alloc_flags_t pmm_flags = UVM_PMM_ALLOC_FLAGS_EVICT;
@@ -1544,25 +1578,28 @@ static NV_STATUS entry_test_maxwell(uvm_gpu_t *gpu)
     uvm_mmu_page_table_alloc_t alloc_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x9999999000LL);
     uvm_mmu_page_table_alloc_t alloc_vid = fake_table_alloc(UVM_APERTURE_VID, 0x1BBBBBB000LL);
     uvm_mmu_mode_hal_t *hal;
+    uvm_page_directory_t dir;
     NvU32 i, j, big_page_size, page_size;
 
+    dir.depth = 0;
+
     for (i = 0; i < ARRAY_SIZE(big_page_sizes); i++) {
         big_page_size = big_page_sizes[i];
         hal = gpu->parent->arch_hal->mmu_mode_hal(big_page_size);
 
         memset(phys_allocs, 0, sizeof(phys_allocs));
 
-        hal->make_pde(&pde_bits, phys_allocs, 0);
+        hal->make_pde(&pde_bits, phys_allocs, &dir, 0);
         TEST_CHECK_RET(pde_bits == 0x0L);
 
         phys_allocs[0] = &alloc_sys;
         phys_allocs[1] = &alloc_vid;
-        hal->make_pde(&pde_bits, phys_allocs, 0);
+        hal->make_pde(&pde_bits, phys_allocs, &dir, 0);
         TEST_CHECK_RET(pde_bits == 0x1BBBBBBD99999992LL);
 
         phys_allocs[0] = &alloc_vid;
         phys_allocs[1] = &alloc_sys;
-        hal->make_pde(&pde_bits, phys_allocs, 0);
+        hal->make_pde(&pde_bits, phys_allocs, &dir, 0);
         TEST_CHECK_RET(pde_bits == 0x9999999E1BBBBBB1LL);
 
         for (j = 0; j <= 2; j++) {
@@ -1632,38 +1669,47 @@ static NV_STATUS entry_test_pascal(uvm_gpu_t *gpu, entry_test_page_size_func ent
     uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
     uvm_mmu_page_table_alloc_t alloc_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x399999999999000LL);
     uvm_mmu_page_table_alloc_t alloc_vid = fake_table_alloc(UVM_APERTURE_VID, 0x1BBBBBB000LL);
+    uvm_page_directory_t dir;
+
     // big versions have [11:8] set as well to test the page table merging
     uvm_mmu_page_table_alloc_t alloc_big_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x399999999999900LL);
     uvm_mmu_page_table_alloc_t alloc_big_vid = fake_table_alloc(UVM_APERTURE_VID, 0x1BBBBBBB00LL);
 
     uvm_mmu_mode_hal_t *hal = gpu->parent->arch_hal->mmu_mode_hal(UVM_PAGE_SIZE_64K);
 
+    dir.index_in_parent = 0;
+    dir.host_parent = NULL;
+    dir.depth = 0;
+
     // Make sure cleared PDEs work as expected
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0);
 
     memset(pde_bits, 0xFF, sizeof(pde_bits));
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    dir.depth = 3;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0 && pde_bits[1] == 0);
 
     // Sys and vidmem PDEs
     phys_allocs[0] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    dir.depth = 0;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x3999999999990C);
 
     phys_allocs[0] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x1BBBBBB0A);
 
     // Dual PDEs
     phys_allocs[0] = &alloc_big_sys;
     phys_allocs[1] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    dir.depth = 3;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x3999999999999C && pde_bits[1] == 0x1BBBBBB0A);
 
     phys_allocs[0] = &alloc_big_vid;
     phys_allocs[1] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x1BBBBBBBA && pde_bits[1] == 0x3999999999990C);
 
     // uncached, i.e., the sysmem data is not cached in GPU's L2 cache. Clear
@@ -1719,6 +1765,7 @@ static NV_STATUS entry_test_volta(uvm_gpu_t *gpu, entry_test_page_size_func entr
     uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
     uvm_mmu_page_table_alloc_t alloc_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x399999999999000LL);
     uvm_mmu_page_table_alloc_t alloc_vid = fake_table_alloc(UVM_APERTURE_VID, 0x1BBBBBB000LL);
+    uvm_page_directory_t dir;
 
     // big versions have [11:8] set as well to test the page table merging
     uvm_mmu_page_table_alloc_t alloc_big_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x399999999999900LL);
@@ -1726,37 +1773,45 @@ static NV_STATUS entry_test_volta(uvm_gpu_t *gpu, entry_test_page_size_func entr
 
     uvm_mmu_mode_hal_t *hal = gpu->parent->arch_hal->mmu_mode_hal(UVM_PAGE_SIZE_64K);
 
+    dir.index_in_parent = 0;
+    dir.host_parent = NULL;
+    dir.depth = 0;
+
     // Make sure cleared PDEs work as expected
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0);
 
     memset(pde_bits, 0xFF, sizeof(pde_bits));
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    dir.depth = 3;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0 && pde_bits[1] == 0);
 
     // Sys and vidmem PDEs
     phys_allocs[0] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    dir.depth = 0;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x3999999999990C);
 
     phys_allocs[0] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 0);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x1BBBBBB0A);
 
     // Dual PDEs
     phys_allocs[0] = &alloc_big_sys;
     phys_allocs[1] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    dir.depth = 3;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x3999999999999C && pde_bits[1] == 0x1BBBBBB0A);
 
     phys_allocs[0] = &alloc_big_vid;
     phys_allocs[1] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 3);
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     TEST_CHECK_RET(pde_bits[0] == 0x1BBBBBBBA && pde_bits[1] == 0x3999999999990C);
 
     // NO_ATS PDE1 (depth 2)
     phys_allocs[0] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 2);
+    dir.depth = 2;
+    hal->make_pde(pde_bits, phys_allocs, &dir, 0);
     if (g_uvm_global.ats.enabled)
         TEST_CHECK_RET(pde_bits[0] == 0x1BBBBBB2A);
     else
@@ -1791,104 +1846,203 @@ static NV_STATUS entry_test_ampere(uvm_gpu_t *gpu, entry_test_page_size_func ent
 
 static NV_STATUS entry_test_hopper(uvm_gpu_t *gpu, entry_test_page_size_func entry_test_page_size)
 {
+    NV_STATUS status = NV_OK;
     NvU32 page_sizes[MAX_NUM_PAGE_SIZES];
     NvU64 pde_bits[2];
+    uvm_page_directory_t *dirs[5];
     size_t i, num_page_sizes;
     uvm_mmu_page_table_alloc_t *phys_allocs[2] = {NULL, NULL};
     uvm_mmu_page_table_alloc_t alloc_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x9999999999000LL);
     uvm_mmu_page_table_alloc_t alloc_vid = fake_table_alloc(UVM_APERTURE_VID, 0xBBBBBBB000LL);
 
-    // big versions have [11:8] set as well to test the page table merging
+    // Big versions have [11:8] set as well to test the page table merging
     uvm_mmu_page_table_alloc_t alloc_big_sys = fake_table_alloc(UVM_APERTURE_SYS, 0x9999999999900LL);
     uvm_mmu_page_table_alloc_t alloc_big_vid = fake_table_alloc(UVM_APERTURE_VID, 0xBBBBBBBB00LL);
 
     uvm_mmu_mode_hal_t *hal = gpu->parent->arch_hal->mmu_mode_hal(UVM_PAGE_SIZE_64K);
 
-    // Make sure cleared PDEs work as expected
-    hal->make_pde(pde_bits, phys_allocs, 0);
-    TEST_CHECK_RET(pde_bits[0] == 0);
+    memset(dirs, 0, sizeof(dirs));
+    // Fake directory tree.
+    for (i = 0; i < ARRAY_SIZE(dirs); i++) {
+        dirs[i] = uvm_kvmalloc_zero(sizeof(uvm_page_directory_t) + sizeof(dirs[i]->entries[0]) * 512);
+        TEST_CHECK_GOTO(dirs[i] != NULL, cleanup);
+
+        dirs[i]->depth = i;
+        dirs[i]->index_in_parent = 0;
+
+        if (i == 0)
+            dirs[i]->host_parent = NULL;
+        else
+            dirs[i]->host_parent = dirs[i - 1];
+    }
+
+    // Make sure cleared PDEs work as expected.
+    hal->make_pde(pde_bits, phys_allocs, dirs[0], 0);
+    TEST_CHECK_GOTO(pde_bits[0] == 0, cleanup);
 
     // Cleared PDEs work as expected for big and small PDEs.
     memset(pde_bits, 0xFF, sizeof(pde_bits));
-    hal->make_pde(pde_bits, phys_allocs, 4);
-    TEST_CHECK_RET(pde_bits[0] == 0 && pde_bits[1] == 0);
+    hal->make_pde(pde_bits, phys_allocs, dirs[4], 0);
+    TEST_CHECK_GOTO(pde_bits[0] == 0 && pde_bits[1] == 0, cleanup);
 
     // Sys and vidmem PDEs, uncached ATS allowed.
     phys_allocs[0] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 0);
-    TEST_CHECK_RET(pde_bits[0] == 0x999999999900C);
+    hal->make_pde(pde_bits, phys_allocs, dirs[0], 0);
+    TEST_CHECK_GOTO(pde_bits[0] == 0x999999999900C, cleanup);
 
     phys_allocs[0] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 0);
-    TEST_CHECK_RET(pde_bits[0] == 0xBBBBBBB00A);
+    hal->make_pde(pde_bits, phys_allocs, dirs[0], 0);
+    TEST_CHECK_GOTO(pde_bits[0] == 0xBBBBBBB00A, cleanup);
 
-    // Dual PDEs, uncached.
+    // Dual PDEs, uncached. We don't use child_dir in the depth 4 checks because
+    // our policy decides the PDE's PCF without using it.
     phys_allocs[0] = &alloc_big_sys;
     phys_allocs[1] = &alloc_vid;
-    hal->make_pde(pde_bits, phys_allocs, 4);
-    TEST_CHECK_RET(pde_bits[0] == 0x999999999991C && pde_bits[1] == 0xBBBBBBB01A);
+    hal->make_pde(pde_bits, phys_allocs, dirs[4], 0);
+    if (g_uvm_global.ats.enabled)
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999991C && pde_bits[1] == 0xBBBBBBB01A, cleanup);
+    else
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999990C && pde_bits[1] == 0xBBBBBBB00A, cleanup);
 
     phys_allocs[0] = &alloc_big_vid;
     phys_allocs[1] = &alloc_sys;
-    hal->make_pde(pde_bits, phys_allocs, 4);
-    TEST_CHECK_RET(pde_bits[0] == 0xBBBBBBBB1A && pde_bits[1] == 0x999999999901C);
+    hal->make_pde(pde_bits, phys_allocs, dirs[4], 0);
+    if (g_uvm_global.ats.enabled)
+        TEST_CHECK_GOTO(pde_bits[0] == 0xBBBBBBBB1A && pde_bits[1] == 0x999999999901C, cleanup);
+    else
+        TEST_CHECK_GOTO(pde_bits[0] == 0xBBBBBBBB0A && pde_bits[1] == 0x999999999900C, cleanup);
+
+    // We only need to test make_pde() on ATS when the CPU VA width < GPU's.
+    if (g_uvm_global.ats.enabled && uvm_cpu_num_va_bits() < hal->num_va_bits()) {
+        phys_allocs[0] = &alloc_sys;
+
+        dirs[1]->index_in_parent = 0;
+        hal->make_pde(pde_bits, phys_allocs, dirs[0], 0);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999900C, cleanup);
+
+        dirs[2]->index_in_parent = 0;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 0);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[2]->index_in_parent = 1;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 1);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[2]->index_in_parent = 2;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 2);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[2]->index_in_parent = 511;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 511);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[1]->index_in_parent = 1;
+        hal->make_pde(pde_bits, phys_allocs, dirs[0], 1);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999900C, cleanup);
+
+        dirs[2]->index_in_parent = 0;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 0);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[2]->index_in_parent = 509;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 509);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        dirs[2]->index_in_parent = 510;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 510);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x999999999901C, cleanup);
+
+        phys_allocs[0] = NULL;
+
+        dirs[1]->index_in_parent = 0;
+        hal->make_pde(pde_bits, phys_allocs, dirs[0], 0);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x0, cleanup);
+
+        dirs[2]->index_in_parent = 0;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 0);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x0, cleanup);
+
+        dirs[2]->index_in_parent = 2;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 2);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x10, cleanup);
+
+        dirs[1]->index_in_parent = 1;
+        dirs[2]->index_in_parent = 509;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 509);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x10, cleanup);
+
+        dirs[2]->index_in_parent = 510;
+        hal->make_pde(pde_bits, phys_allocs, dirs[1], 510);
+        TEST_CHECK_GOTO(pde_bits[0] == 0x0, cleanup);
+    }
 
     // uncached, i.e., the sysmem data is not cached in GPU's L2 cache, and
     // access counters disabled.
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_SYS,
-                                 0x9999999999000LL,
-                                 UVM_PROT_READ_WRITE_ATOMIC,
-                                 UVM_MMU_PTE_FLAGS_ACCESS_COUNTERS_DISABLED) == 0x999999999968D);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_SYS,
+                                  0x9999999999000LL,
+                                  UVM_PROT_READ_WRITE_ATOMIC,
+                                  UVM_MMU_PTE_FLAGS_ACCESS_COUNTERS_DISABLED) == 0x999999999968D,
+                    cleanup);
 
     // change to cached.
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_SYS,
-                                 0x9999999999000LL,
-                                 UVM_PROT_READ_WRITE_ATOMIC,
-                                 UVM_MMU_PTE_FLAGS_CACHED | UVM_MMU_PTE_FLAGS_ACCESS_COUNTERS_DISABLED) ==
-                   0x9999999999685);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_SYS,
+                                  0x9999999999000LL,
+                                  UVM_PROT_READ_WRITE_ATOMIC,
+                                  UVM_MMU_PTE_FLAGS_CACHED | UVM_MMU_PTE_FLAGS_ACCESS_COUNTERS_DISABLED) ==
+                                  0x9999999999685,
+                    cleanup);
 
     // enable access counters.
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_SYS,
-                                 0x9999999999000LL,
-                                 UVM_PROT_READ_WRITE_ATOMIC,
-                                 UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999605);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_SYS,
+                                  0x9999999999000LL,
+                                  UVM_PROT_READ_WRITE_ATOMIC,
+                                  UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999605,
+                    cleanup);
 
     // remove atomic
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_SYS,
-                                 0x9999999999000LL,
-                                 UVM_PROT_READ_WRITE,
-                                 UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999645);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_SYS,
+                                  0x9999999999000LL,
+                                  UVM_PROT_READ_WRITE,
+                                  UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999645,
+                    cleanup);
 
     // read only
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_SYS,
-                                 0x9999999999000LL,
-                                 UVM_PROT_READ_ONLY,
-                                 UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999665);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_SYS,
+                                  0x9999999999000LL,
+                                  UVM_PROT_READ_ONLY,
+                                  UVM_MMU_PTE_FLAGS_CACHED) == 0x9999999999665,
+                    cleanup);
 
     // local video
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_VID,
-                                 0xBBBBBBB000LL,
-                                 UVM_PROT_READ_ONLY,
-                                 UVM_MMU_PTE_FLAGS_CACHED) == 0xBBBBBBB661);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_VID,
+                                  0xBBBBBBB000LL,
+                                  UVM_PROT_READ_ONLY,
+                                  UVM_MMU_PTE_FLAGS_CACHED) == 0xBBBBBBB661,
+                    cleanup);
 
     // peer 1
-    TEST_CHECK_RET(hal->make_pte(UVM_APERTURE_PEER_1,
-                                 0xBBBBBBB000LL,
-                                 UVM_PROT_READ_ONLY,
-                                 UVM_MMU_PTE_FLAGS_CACHED) == 0x200000BBBBBBB663);
+    TEST_CHECK_GOTO(hal->make_pte(UVM_APERTURE_PEER_1,
+                                  0xBBBBBBB000LL,
+                                  UVM_PROT_READ_ONLY,
+                                  UVM_MMU_PTE_FLAGS_CACHED) == 0x200000BBBBBBB663,
+                    cleanup);
 
     // sparse
-    TEST_CHECK_RET(hal->make_sparse_pte() == 0x8);
+    TEST_CHECK_GOTO(hal->make_sparse_pte() == 0x8, cleanup);
 
     // sked reflected
-    TEST_CHECK_RET(hal->make_sked_reflected_pte() == 0xF09);
+    TEST_CHECK_GOTO(hal->make_sked_reflected_pte() == 0xF09, cleanup);
 
     num_page_sizes = get_page_sizes(gpu, page_sizes);
 
     for (i = 0; i < num_page_sizes; i++)
-        TEST_NV_CHECK_RET(entry_test_page_size(gpu, page_sizes[i]));
+        TEST_NV_CHECK_GOTO(entry_test_page_size(gpu, page_sizes[i]), cleanup);
 
-    return NV_OK;
+cleanup:
+    for (i = 0; i < ARRAY_SIZE(dirs); i++)
+        uvm_kvfree(dirs[i]);
+
+    return status;
 }
 
 static NV_STATUS alloc_4k_maxwell(uvm_gpu_t *gpu)
@@ -2303,7 +2457,8 @@ NV_STATUS uvm_test_page_tree(UVM_TEST_PAGE_TREE_PARAMS *params, struct file *fil
     gpu->parent = parent_gpu;
 
     // At least test_tlb_invalidates() relies on global state
-    // (g_tlb_invalidate_*) so make sure only one test instance can run at a time.
+    // (g_tlb_invalidate_*) so make sure only one test instance can run at a
+    // time.
     uvm_mutex_lock(&g_uvm_global.global_lock);
 
     // Allocate the fake TLB tracking state. Notably tests still need to enable
@@ -2311,7 +2466,13 @@ NV_STATUS uvm_test_page_tree(UVM_TEST_PAGE_TREE_PARAMS *params, struct file *fil
     // calls.
     TEST_NV_CHECK_GOTO(fake_tlb_invals_alloc(), done);
 
-    TEST_NV_CHECK_GOTO(maxwell_test_page_tree(gpu), done);
+    // We prevent the maxwell_test_page_tree test from running on ATS-enabled
+    // systems. On "fake" Maxwell-based ATS systems pde_fill() may push more
+    // methods than what we support in UVM. Specifically, on
+    // uvm_page_tree_init() which eventually calls phys_mem_init(). On Maxwell,
+    // upper PDE levels have more than 512 entries.
+    if (!g_uvm_global.ats.enabled)
+        TEST_NV_CHECK_GOTO(maxwell_test_page_tree(gpu), done);
     TEST_NV_CHECK_GOTO(pascal_test_page_tree(gpu), done);
     TEST_NV_CHECK_GOTO(volta_test_page_tree(gpu), done);
     TEST_NV_CHECK_GOTO(ampere_test_page_tree(gpu), done);

kernel-open/nvidia-uvm/uvm_pascal.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2016-2020 NVIDIA Corporation
+    Copyright (c) 2016-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -100,4 +100,6 @@ void uvm_hal_pascal_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     parent_gpu->smc.supported = false;
 
     parent_gpu->plc_supported = false;
+
+    parent_gpu->no_ats_range_required = false;
 }

kernel-open/nvidia-uvm/uvm_pascal_mmu.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2020 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -140,11 +140,18 @@ static NvU64 small_half_pde_pascal(uvm_mmu_page_table_alloc_t *phys_alloc)
     return pde_bits;
 }
 
-static void make_pde_pascal(void *entry, uvm_mmu_page_table_alloc_t **phys_allocs, NvU32 depth)
+static void make_pde_pascal(void *entry,
+                            uvm_mmu_page_table_alloc_t **phys_allocs,
+                            uvm_page_directory_t *dir,
+                            NvU32 child_index)
 {
-    NvU32 entry_count = entries_per_index_pascal(depth);
+    NvU32 entry_count;
     NvU64 *entry_bits = (NvU64 *)entry;
 
+    UVM_ASSERT(dir);
+
+    entry_count = entries_per_index_pascal(dir->depth);
+
     if (entry_count == 1) {
         *entry_bits = single_pde_pascal(*phys_allocs);
     }
@@ -152,7 +159,8 @@ static void make_pde_pascal(void *entry, uvm_mmu_page_table_alloc_t **phys_alloc
         entry_bits[MMU_BIG] = big_half_pde_pascal(phys_allocs[MMU_BIG]);
         entry_bits[MMU_SMALL] = small_half_pde_pascal(phys_allocs[MMU_SMALL]);
 
-        // This entry applies to the whole dual PDE but is stored in the lower bits
+        // This entry applies to the whole dual PDE but is stored in the lower
+        // bits.
         entry_bits[MMU_BIG] |= HWCONST64(_MMU_VER2, DUAL_PDE, IS_PDE, TRUE);
     }
     else {

kernel-open/nvidia-uvm/uvm_perf_prefetch.c

@@ -218,57 +218,11 @@ static void grow_fault_granularity(uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
     }
 }
 
-// Within a block we only allow prefetching to a single processor. Therefore,
-// if two processors are accessing non-overlapping regions within the same
-// block they won't benefit from prefetching.
-//
-// TODO: Bug 1778034: [uvm] Explore prefetching to different processors within
-// a VA block.
-static NvU32 uvm_perf_prefetch_prenotify_fault_migrations(uvm_va_block_t *va_block,
-                                                          uvm_va_block_context_t *va_block_context,
-                                                          uvm_processor_id_t new_residency,
-                                                          const uvm_page_mask_t *faulted_pages,
-                                                          uvm_va_block_region_t faulted_region,
-                                                          uvm_page_mask_t *prefetch_pages,
-                                                          uvm_perf_prefetch_bitmap_tree_t *bitmap_tree)
+static void init_bitmap_tree_from_region(uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                         uvm_va_block_region_t max_prefetch_region,
+                                         const uvm_page_mask_t *resident_mask,
+                                         const uvm_page_mask_t *faulted_pages)
 {
-    uvm_page_index_t page_index;
-    const uvm_page_mask_t *resident_mask = NULL;
-    const uvm_page_mask_t *thrashing_pages = NULL;
-    uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
-    const uvm_va_policy_t *policy = va_block_context->policy;
-    uvm_va_block_region_t max_prefetch_region;
-    NvU32 big_page_size;
-    uvm_va_block_region_t big_pages_region;
-
-    if (!uvm_id_equal(va_block->prefetch_info.last_migration_proc_id, new_residency)) {
-        va_block->prefetch_info.last_migration_proc_id = new_residency;
-        va_block->prefetch_info.fault_migrations_to_last_proc = 0;
-    }
-
-    // Compute the expanded region that prefetching is allowed from.
-    if (uvm_va_block_is_hmm(va_block)) {
-        max_prefetch_region = uvm_hmm_get_prefetch_region(va_block,
-                                                          va_block_context,
-                                                          uvm_va_block_region_start(va_block, faulted_region));
-    }
-    else {
-        max_prefetch_region = uvm_va_block_region_from_block(va_block);
-    }
-
-    uvm_page_mask_zero(prefetch_pages);
-
-    if (UVM_ID_IS_CPU(new_residency) || va_block->gpus[uvm_id_gpu_index(new_residency)] != NULL)
-        resident_mask = uvm_va_block_resident_mask_get(va_block, new_residency);
-
-    // If this is a first-touch fault and the destination processor is the
-    // preferred location, populate the whole max_prefetch_region.
-    if (uvm_processor_mask_empty(&va_block->resident) &&
-        uvm_id_equal(new_residency, policy->preferred_location)) {
-        uvm_page_mask_region_fill(prefetch_pages, max_prefetch_region);
-        goto done;
-    }
-
     if (resident_mask)
         uvm_page_mask_or(&bitmap_tree->pages, resident_mask, faulted_pages);
     else
@@ -277,6 +231,29 @@ static NvU32 uvm_perf_prefetch_prenotify_fault_migrations(uvm_va_block_t *va_blo
     // If we are using a subregion of the va_block, align bitmap_tree
     uvm_page_mask_shift_right(&bitmap_tree->pages, &bitmap_tree->pages, max_prefetch_region.first);
 
+    bitmap_tree->offset = 0;
+    bitmap_tree->leaf_count = uvm_va_block_region_num_pages(max_prefetch_region);
+    bitmap_tree->level_count = ilog2(roundup_pow_of_two(bitmap_tree->leaf_count)) + 1;
+}
+
+static void update_bitmap_tree_from_va_block(uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                             uvm_va_block_t *va_block,
+                                             uvm_va_block_context_t *va_block_context,
+                                             uvm_processor_id_t new_residency,
+                                             const uvm_page_mask_t *faulted_pages,
+                                             uvm_va_block_region_t max_prefetch_region)
+
+{
+    NvU32 big_page_size;
+    uvm_va_block_region_t big_pages_region;
+    uvm_va_space_t *va_space;
+    const uvm_page_mask_t *thrashing_pages;
+
+    UVM_ASSERT(va_block);
+    UVM_ASSERT(va_block_context);
+
+    va_space = uvm_va_block_get_va_space(va_block);
+
     // Get the big page size for the new residency.
     // Assume 64K size if the new residency is the CPU or no GPU va space is
     // registered in the current process for this GPU.
@@ -302,13 +279,9 @@ static NvU32 uvm_perf_prefetch_prenotify_fault_migrations(uvm_va_block_t *va_blo
         UVM_ASSERT(bitmap_tree->leaf_count <= PAGES_PER_UVM_VA_BLOCK);
 
         uvm_page_mask_shift_left(&bitmap_tree->pages, &bitmap_tree->pages, bitmap_tree->offset);
-    }
-    else {
-        bitmap_tree->offset = 0;
-        bitmap_tree->leaf_count = uvm_va_block_region_num_pages(max_prefetch_region);
-    }
 
-    bitmap_tree->level_count = ilog2(roundup_pow_of_two(bitmap_tree->leaf_count)) + 1;
+        bitmap_tree->level_count = ilog2(roundup_pow_of_two(bitmap_tree->leaf_count)) + 1;
+    }
 
     thrashing_pages = uvm_perf_thrashing_get_thrashing_pages(va_block);
 
@@ -320,25 +293,99 @@ static NvU32 uvm_perf_prefetch_prenotify_fault_migrations(uvm_va_block_t *va_blo
                            max_prefetch_region,
                            faulted_pages,
                            thrashing_pages);
+}
 
-    // Do not compute prefetch regions with faults on pages that are thrashing
-    if (thrashing_pages)
-        uvm_page_mask_andnot(&va_block_context->scratch_page_mask, faulted_pages, thrashing_pages);
-    else
-        uvm_page_mask_copy(&va_block_context->scratch_page_mask, faulted_pages);
+static void compute_prefetch_mask(uvm_va_block_region_t faulted_region,
+                                  uvm_va_block_region_t max_prefetch_region,
+                                  uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                  const uvm_page_mask_t *faulted_pages,
+                                  uvm_page_mask_t *out_prefetch_mask)
+{
+    uvm_page_index_t page_index;
 
-    // Update the tree using the scratch mask to compute the pages to prefetch
-    for_each_va_block_page_in_region_mask(page_index, &va_block_context->scratch_page_mask, faulted_region) {
+    uvm_page_mask_zero(out_prefetch_mask);
+
+    // Update the tree using the faulted mask to compute the pages to prefetch.
+    for_each_va_block_page_in_region_mask(page_index, faulted_pages, faulted_region) {
         uvm_va_block_region_t region = compute_prefetch_region(page_index, bitmap_tree, max_prefetch_region);
 
-        uvm_page_mask_region_fill(prefetch_pages, region);
+        uvm_page_mask_region_fill(out_prefetch_mask, region);
 
         // Early out if we have already prefetched until the end of the VA block
         if (region.outer == max_prefetch_region.outer)
             break;
     }
+}
+
+// Within a block we only allow prefetching to a single processor. Therefore,
+// if two processors are accessing non-overlapping regions within the same
+// block they won't benefit from prefetching.
+//
+// TODO: Bug 1778034: [uvm] Explore prefetching to different processors within
+// a VA block.
+static NvU32 uvm_perf_prefetch_prenotify_fault_migrations(uvm_va_block_t *va_block,
+                                                          uvm_va_block_context_t *va_block_context,
+                                                          uvm_processor_id_t new_residency,
+                                                          const uvm_page_mask_t *faulted_pages,
+                                                          uvm_va_block_region_t faulted_region,
+                                                          uvm_page_mask_t *prefetch_pages,
+                                                          uvm_perf_prefetch_bitmap_tree_t *bitmap_tree)
+{
+    const uvm_page_mask_t *resident_mask = NULL;
+    const uvm_va_policy_t *policy = uvm_va_policy_get_region(va_block, faulted_region);
+    uvm_va_block_region_t max_prefetch_region;
+    const uvm_page_mask_t *thrashing_pages = uvm_perf_thrashing_get_thrashing_pages(va_block);
+
+    if (!uvm_id_equal(va_block->prefetch_info.last_migration_proc_id, new_residency)) {
+        va_block->prefetch_info.last_migration_proc_id = new_residency;
+        va_block->prefetch_info.fault_migrations_to_last_proc = 0;
+    }
+
+    // Compute the expanded region that prefetching is allowed from.
+    if (uvm_va_block_is_hmm(va_block)) {
+        max_prefetch_region = uvm_hmm_get_prefetch_region(va_block,
+                                                          va_block_context->hmm.vma,
+                                                          policy,
+                                                          uvm_va_block_region_start(va_block, faulted_region));
+    }
+    else {
+        max_prefetch_region = uvm_va_block_region_from_block(va_block);
+    }
+
+    uvm_page_mask_zero(prefetch_pages);
+
+    if (UVM_ID_IS_CPU(new_residency) || va_block->gpus[uvm_id_gpu_index(new_residency)] != NULL)
+        resident_mask = uvm_va_block_resident_mask_get(va_block, new_residency);
+
+    // If this is a first-touch fault and the destination processor is the
+    // preferred location, populate the whole max_prefetch_region.
+    if (uvm_processor_mask_empty(&va_block->resident) &&
+        uvm_id_equal(new_residency, policy->preferred_location)) {
+        uvm_page_mask_region_fill(prefetch_pages, max_prefetch_region);
+    }
+    else {
+        init_bitmap_tree_from_region(bitmap_tree, max_prefetch_region, resident_mask, faulted_pages);
+
+        update_bitmap_tree_from_va_block(bitmap_tree,
+                                         va_block,
+                                         va_block_context,
+                                         new_residency,
+                                         faulted_pages,
+                                         max_prefetch_region);
+
+        // Do not compute prefetch regions with faults on pages that are thrashing
+        if (thrashing_pages)
+            uvm_page_mask_andnot(&va_block_context->scratch_page_mask, faulted_pages, thrashing_pages);
+        else
+            uvm_page_mask_copy(&va_block_context->scratch_page_mask, faulted_pages);
+
+        compute_prefetch_mask(faulted_region,
+                              max_prefetch_region,
+                              bitmap_tree,
+                              &va_block_context->scratch_page_mask,
+                              prefetch_pages);
+    }
 
-done:
     // Do not prefetch pages that are going to be migrated/populated due to a
     // fault
     uvm_page_mask_andnot(prefetch_pages, prefetch_pages, faulted_pages);
@@ -364,31 +411,58 @@ done:
     return uvm_page_mask_weight(prefetch_pages);
 }
 
-void uvm_perf_prefetch_get_hint(uvm_va_block_t *va_block,
-                                uvm_va_block_context_t *va_block_context,
-                                uvm_processor_id_t new_residency,
-                                const uvm_page_mask_t *faulted_pages,
-                                uvm_va_block_region_t faulted_region,
-                                uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
-                                uvm_perf_prefetch_hint_t *out_hint)
+bool uvm_perf_prefetch_enabled(uvm_va_space_t *va_space)
+{
+    if (!g_uvm_perf_prefetch_enable)
+        return false;
+
+    UVM_ASSERT(va_space);
+
+    return va_space->test.page_prefetch_enabled;
+}
+
+void uvm_perf_prefetch_compute_ats(uvm_va_space_t *va_space,
+                                   const uvm_page_mask_t *faulted_pages,
+                                   uvm_va_block_region_t faulted_region,
+                                   uvm_va_block_region_t max_prefetch_region,
+                                   const uvm_page_mask_t *residency_mask,
+                                   uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                   uvm_page_mask_t *out_prefetch_mask)
+{
+    UVM_ASSERT(faulted_pages);
+    UVM_ASSERT(bitmap_tree);
+    UVM_ASSERT(out_prefetch_mask);
+
+    uvm_page_mask_zero(out_prefetch_mask);
+
+    if (!uvm_perf_prefetch_enabled(va_space))
+        return;
+
+    init_bitmap_tree_from_region(bitmap_tree, max_prefetch_region, residency_mask, faulted_pages);
+
+    compute_prefetch_mask(faulted_region, max_prefetch_region, bitmap_tree, faulted_pages, out_prefetch_mask);
+}
+
+void uvm_perf_prefetch_get_hint_va_block(uvm_va_block_t *va_block,
+                                         uvm_va_block_context_t *va_block_context,
+                                         uvm_processor_id_t new_residency,
+                                         const uvm_page_mask_t *faulted_pages,
+                                         uvm_va_block_region_t faulted_region,
+                                         uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                         uvm_perf_prefetch_hint_t *out_hint)
 {
-    const uvm_va_policy_t *policy = va_block_context->policy;
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
     uvm_page_mask_t *prefetch_pages = &out_hint->prefetch_pages_mask;
     NvU32 pending_prefetch_pages;
 
     uvm_assert_rwsem_locked(&va_space->lock);
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, policy, faulted_region));
-    UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block, va_block_context, faulted_region));
+    UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block, va_block_context->hmm.vma, faulted_region));
 
     out_hint->residency = UVM_ID_INVALID;
     uvm_page_mask_zero(prefetch_pages);
 
-    if (!g_uvm_perf_prefetch_enable)
-        return;
-
-    if (!va_space->test.page_prefetch_enabled)
+    if (!uvm_perf_prefetch_enabled(va_space))
         return;
 
     pending_prefetch_pages = uvm_perf_prefetch_prenotify_fault_migrations(va_block,

kernel-open/nvidia-uvm/uvm_perf_prefetch.h

@@ -61,21 +61,41 @@ typedef struct
 // Global initialization function (no clean up needed).
 NV_STATUS uvm_perf_prefetch_init(void);
 
+// Returns whether prefetching is enabled in the VA space.
+// va_space cannot be NULL.
+bool uvm_perf_prefetch_enabled(uvm_va_space_t *va_space);
+
+// Return the prefetch mask with the pages that may be prefetched in a ATS
+// block. ATS block is a system allocated memory block with base aligned to
+// UVM_VA_BLOCK_SIZE and a maximum size of UVM_VA_BLOCK_SIZE. The faulted_pages
+// mask and faulted_region are the pages being faulted on the given residency.
+//
+// Only residency_mask can be NULL.
+//
+// Locking: The caller must hold the va_space lock.
+void uvm_perf_prefetch_compute_ats(uvm_va_space_t *va_space,
+                                   const uvm_page_mask_t *faulted_pages,
+                                   uvm_va_block_region_t faulted_region,
+                                   uvm_va_block_region_t max_prefetch_region,
+                                   const uvm_page_mask_t *residency_mask,
+                                   uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                   uvm_page_mask_t *out_prefetch_mask);
+
 // Return a hint with the pages that may be prefetched in the block.
 // The faulted_pages mask and faulted_region are the pages being migrated to
 // the given residency.
-// va_block_context must not be NULL, va_block_context->policy must be valid,
-// and if the va_block is a HMM block, va_block_context->hmm.vma must be valid
-// which also means the va_block_context->mm is not NULL, retained, and locked
-// for at least read.
+// va_block_context must not be NULL, and if the va_block is a HMM
+// block, va_block_context->hmm.vma must be valid which also means the
+// va_block_context->mm is not NULL, retained, and locked for at least
+// read.
 // Locking: The caller must hold the va_space lock and va_block lock.
-void uvm_perf_prefetch_get_hint(uvm_va_block_t *va_block,
-                                uvm_va_block_context_t *va_block_context,
-                                uvm_processor_id_t new_residency,
-                                const uvm_page_mask_t *faulted_pages,
-                                uvm_va_block_region_t faulted_region,
-                                uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
-                                uvm_perf_prefetch_hint_t *out_hint);
+void uvm_perf_prefetch_get_hint_va_block(uvm_va_block_t *va_block,
+                                         uvm_va_block_context_t *va_block_context,
+                                         uvm_processor_id_t new_residency,
+                                         const uvm_page_mask_t *faulted_pages,
+                                         uvm_va_block_region_t faulted_region,
+                                         uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
+                                         uvm_perf_prefetch_hint_t *out_hint);
 
 void uvm_perf_prefetch_bitmap_tree_iter_init(const uvm_perf_prefetch_bitmap_tree_t *bitmap_tree,
                                              uvm_page_index_t page_index,

kernel-open/nvidia-uvm/uvm_perf_thrashing.c

@@ -1095,7 +1095,7 @@ static NV_STATUS unmap_remote_pinned_pages(uvm_va_block_t *va_block,
     NV_STATUS tracker_status;
     uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
     uvm_processor_id_t processor_id;
-    const uvm_va_policy_t *policy = va_block_context->policy;
+    const uvm_va_policy_t *policy = uvm_va_policy_get(va_block, uvm_va_block_region_start(va_block, region));
 
     uvm_assert_mutex_locked(&va_block->lock);
 
@@ -1141,10 +1141,9 @@ NV_STATUS uvm_perf_thrashing_unmap_remote_pinned_pages_all(uvm_va_block_t *va_bl
 {
     block_thrashing_info_t *block_thrashing;
     uvm_processor_mask_t unmap_processors;
-    const uvm_va_policy_t *policy = va_block_context->policy;
+    const uvm_va_policy_t *policy = uvm_va_policy_get_region(va_block, region);
 
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, policy, region));
 
     block_thrashing = thrashing_info_get(va_block);
     if (!block_thrashing || !block_thrashing->pages)
@@ -1867,8 +1866,6 @@ static void thrashing_unpin_pages(struct work_struct *work)
             UVM_ASSERT(uvm_page_mask_test(&block_thrashing->pinned_pages.mask, page_index));
 
             uvm_va_block_context_init(va_block_context, NULL);
-            va_block_context->policy =
-                uvm_va_policy_get(va_block, uvm_va_block_cpu_page_address(va_block, page_index));
 
             uvm_perf_thrashing_unmap_remote_pinned_pages_all(va_block,
                                                              va_block_context,
@@ -2123,8 +2120,6 @@ NV_STATUS uvm_test_set_page_thrashing_policy(UVM_TEST_SET_PAGE_THRASHING_POLICY_
                 uvm_va_block_region_t va_block_region = uvm_va_block_region_from_block(va_block);
                 uvm_va_block_context_t *block_context = uvm_va_space_block_context(va_space, NULL);
 
-                block_context->policy = uvm_va_range_get_policy(va_range);
-
                 uvm_mutex_lock(&va_block->lock);
 
                 // Unmap may split PTEs and require a retry. Needs to be called

kernel-open/nvidia-uvm/uvm_perf_thrashing.h

@@ -103,11 +103,11 @@ void uvm_perf_thrashing_unload(uvm_va_space_t *va_space);
 // Destroy the thrashing detection struct for the given block.
 void uvm_perf_thrashing_info_destroy(uvm_va_block_t *va_block);
 
-// Unmap remote mappings from all processors on the pinned pages
-// described by region and block_thrashing->pinned pages.
-// va_block_context must not be NULL and va_block_context->policy must be valid.
-// See the comments for uvm_va_block_check_policy_is_valid() in uvm_va_block.h.
-// Locking: the va_block lock must be held.
+// Unmap remote mappings from all processors on the pinned pages described by
+// region and block_thrashing->pinned pages.  va_block_context must not be NULL
+// and policy for the region must match.  See the comments for
+// uvm_va_block_check_policy_is_valid() in uvm_va_block.h.  Locking: the
+// va_block lock must be held.
 NV_STATUS uvm_perf_thrashing_unmap_remote_pinned_pages_all(uvm_va_block_t *va_block,
                                                            uvm_va_block_context_t *va_block_context,
                                                            uvm_va_block_region_t region);

kernel-open/nvidia-uvm/uvm_pmm_gpu.c

@@ -221,7 +221,7 @@ struct uvm_pmm_gpu_chunk_suballoc_struct
     // Array of all child subchunks
     // TODO: Bug 1765461: Can the array be inlined? It could save the parent
     //       pointer.
-    uvm_gpu_chunk_t *subchunks[0];
+    uvm_gpu_chunk_t *subchunks[];
 };
 
 typedef enum
@@ -3314,7 +3314,7 @@ NvU32 uvm_pmm_gpu_phys_to_virt(uvm_pmm_gpu_t *pmm, NvU64 phys_addr, NvU64 region
 
 static uvm_pmm_gpu_t *devmem_page_to_pmm(struct page *page)
 {
-    return container_of(page->pgmap, uvm_pmm_gpu_t, devmem.pagemap);
+    return container_of(page_pgmap(page), uvm_pmm_gpu_t, devmem.pagemap);
 }
 
 static uvm_gpu_chunk_t *devmem_page_to_chunk_locked(struct page *page)
@@ -3820,18 +3820,11 @@ NV_STATUS uvm_test_evict_chunk(UVM_TEST_EVICT_CHUNK_PARAMS *params, struct file
     // For virtual mode, look up and retain the block first so that eviction can
     // be started without the VA space lock held.
     if (params->eviction_mode == UvmTestEvictModeVirtual) {
-        uvm_va_block_context_t *block_context;
+        if (mm)
+            status = uvm_va_block_find_create(va_space, params->address, NULL, &block);
+        else
+            status = uvm_va_block_find_create_managed(va_space, params->address, &block);
 
-        block_context = uvm_va_block_context_alloc(mm);
-        if (!block_context) {
-            status = NV_ERR_NO_MEMORY;
-            uvm_va_space_up_read(va_space);
-            uvm_va_space_mm_release_unlock(va_space, mm);
-            goto out;
-        }
-
-        status = uvm_va_block_find_create(va_space, params->address, block_context, &block);
-        uvm_va_block_context_free(block_context);
         if (status != NV_OK) {
             uvm_va_space_up_read(va_space);
             uvm_va_space_mm_or_current_release_unlock(va_space, mm);

kernel-open/nvidia-uvm/uvm_pmm_test.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2015-2022 NVIDIA Corporation
+    Copyright (c) 2015-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -324,7 +324,7 @@ static NV_STATUS gpu_mem_check(uvm_gpu_t *gpu,
 
     // TODO: Bug 3839176: [UVM][HCC][uvm_test] Update tests that assume GPU
     //                     engines can directly access sysmem
-    // Skip this test for now. To enable this test under SEV,
+    // Skip this test for now. To enable this test in Confidential Computing,
     // The GPU->CPU CE copy needs to be updated so it uses encryption when
     // CC is enabled.
     if (uvm_conf_computing_mode_enabled(gpu))
@@ -1223,8 +1223,6 @@ static NV_STATUS test_indirect_peers(uvm_gpu_t *owning_gpu, uvm_gpu_t *accessing
     if (!chunks)
         return NV_ERR_NO_MEMORY;
 
-    UVM_ASSERT(!g_uvm_global.sev_enabled);
-
     TEST_NV_CHECK_GOTO(uvm_mem_alloc_sysmem_and_map_cpu_kernel(UVM_CHUNK_SIZE_MAX, current->mm, &verif_mem), out);
     TEST_NV_CHECK_GOTO(uvm_mem_map_gpu_kernel(verif_mem, owning_gpu), out);
     TEST_NV_CHECK_GOTO(uvm_mem_map_gpu_kernel(verif_mem, accessing_gpu), out);

kernel-open/nvidia-uvm/uvm_policy.c

@@ -160,7 +160,7 @@ static NV_STATUS preferred_location_unmap_remote_pages(uvm_va_block_t *va_block,
     NV_STATUS status = NV_OK;
     NV_STATUS tracker_status;
     uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
-    const uvm_va_policy_t *policy = va_block_context->policy;
+    const uvm_va_policy_t *policy = uvm_va_policy_get_region(va_block, region);
     uvm_processor_id_t preferred_location = policy->preferred_location;
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
     const uvm_page_mask_t *mapped_mask;
@@ -279,6 +279,9 @@ static NV_STATUS preferred_location_set(uvm_va_space_t *va_space,
         return NV_OK;
     }
 
+    if (!mm)
+        return NV_ERR_INVALID_ADDRESS;
+
     return uvm_hmm_set_preferred_location(va_space, preferred_location, base, last_address, out_tracker);
 }
 
@@ -445,7 +448,6 @@ NV_STATUS uvm_va_block_set_accessed_by_locked(uvm_va_block_t *va_block,
     NV_STATUS tracker_status;
 
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, va_block_context->policy, region));
 
     status = uvm_va_block_add_mappings(va_block,
                                        va_block_context,
@@ -467,13 +469,13 @@ NV_STATUS uvm_va_block_set_accessed_by(uvm_va_block_t *va_block,
     uvm_va_block_region_t region = uvm_va_block_region_from_block(va_block);
     NV_STATUS status;
     uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
+    uvm_va_policy_t *policy = uvm_va_range_get_policy(va_block->va_range);
 
     UVM_ASSERT(!uvm_va_block_is_hmm(va_block));
-    UVM_ASSERT(va_block_context->policy == uvm_va_range_get_policy(va_block->va_range));
 
     // Read duplication takes precedence over SetAccessedBy. Do not add mappings
     // if read duplication is enabled.
-    if (uvm_va_policy_is_read_duplicate(va_block_context->policy, va_space))
+    if (uvm_va_policy_is_read_duplicate(policy, va_space))
         return NV_OK;
 
     status = UVM_VA_BLOCK_LOCK_RETRY(va_block,
@@ -592,8 +594,15 @@ static NV_STATUS accessed_by_set(uvm_va_space_t *va_space,
         UVM_ASSERT(va_range_last->node.end >= last_address);
     }
     else {
+        // NULL mm case already filtered by uvm_api_range_type_check()
+        UVM_ASSERT(mm);
         UVM_ASSERT(type == UVM_API_RANGE_TYPE_HMM);
-        status = uvm_hmm_set_accessed_by(va_space, processor_id, set_bit, base, last_address, &local_tracker);
+        status = uvm_hmm_set_accessed_by(va_space,
+                                         processor_id,
+                                         set_bit,
+                                         base,
+                                         last_address,
+                                         &local_tracker);
     }
 
 done:
@@ -656,7 +665,6 @@ NV_STATUS uvm_va_block_set_read_duplication(uvm_va_block_t *va_block,
 
     // TODO: Bug 3660922: need to implement HMM read duplication support.
     UVM_ASSERT(!uvm_va_block_is_hmm(va_block));
-    UVM_ASSERT(va_block_context->policy == uvm_va_range_get_policy(va_block->va_range));
 
     status = UVM_VA_BLOCK_LOCK_RETRY(va_block, &va_block_retry,
                                      va_block_set_read_duplication_locked(va_block,
@@ -675,7 +683,7 @@ static NV_STATUS va_block_unset_read_duplication_locked(uvm_va_block_t *va_block
     uvm_processor_id_t processor_id;
     uvm_va_block_region_t block_region = uvm_va_block_region_from_block(va_block);
     uvm_page_mask_t *break_read_duplication_pages = &va_block_context->caller_page_mask;
-    const uvm_va_policy_t *policy = va_block_context->policy;
+    const uvm_va_policy_t *policy = uvm_va_range_get_policy(va_block->va_range);
     uvm_processor_id_t preferred_location = policy->preferred_location;
     uvm_processor_mask_t accessed_by = policy->accessed_by;
 
@@ -757,7 +765,6 @@ NV_STATUS uvm_va_block_unset_read_duplication(uvm_va_block_t *va_block,
     uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
 
     UVM_ASSERT(!uvm_va_block_is_hmm(va_block));
-    UVM_ASSERT(va_block_context->policy == uvm_va_range_get_policy(va_block->va_range));
 
     // Restore all SetAccessedBy mappings
     status = UVM_VA_BLOCK_LOCK_RETRY(va_block, &va_block_retry,
@@ -915,7 +922,6 @@ static NV_STATUS system_wide_atomics_set(uvm_va_space_t *va_space, const NvProce
             if (va_range->type != UVM_VA_RANGE_TYPE_MANAGED)
                 continue;
 
-            va_block_context->policy = uvm_va_range_get_policy(va_range);
             for_each_va_block_in_va_range(va_range, va_block) {
                 uvm_page_mask_t *non_resident_pages = &va_block_context->caller_page_mask;
 

kernel-open/nvidia-uvm/uvm_push.c

@@ -391,11 +391,13 @@ uvm_gpu_address_t uvm_push_inline_data_end(uvm_push_inline_data_t *data)
         inline_data_address = (NvU64) (uintptr_t)(push->next + 1);
     }
     else {
+        uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
+
         // Offset of the inlined data within the push.
         inline_data_address = (push->next - push->begin + 1) * UVM_METHOD_SIZE;
 
         // Add GPU VA of the push begin
-        inline_data_address += uvm_pushbuffer_get_gpu_va_for_push(channel->pool->manager->pushbuffer, push);
+        inline_data_address += uvm_pushbuffer_get_gpu_va_for_push(pushbuffer, push);
     }
 
     // This will place a noop right before the inline data that was written.
@@ -438,10 +440,8 @@ NvU64 *uvm_push_timestamp(uvm_push_t *push)
 
     if (uvm_channel_is_ce(push->channel))
         gpu->parent->ce_hal->semaphore_timestamp(push, address.address);
-    else if (uvm_channel_is_sec2(push->channel))
-        gpu->parent->sec2_hal->semaphore_timestamp(push, address.address);
     else
-        UVM_ASSERT_MSG(0, "Semaphore release timestamp on an unsupported channel.\n");
+        gpu->parent->sec2_hal->semaphore_timestamp(push, address.address);
 
     return timestamp;
 }

kernel-open/nvidia-uvm/uvm_push.h

@@ -377,11 +377,6 @@ static bool uvm_push_has_space(uvm_push_t *push, NvU32 free_space)
 NV_STATUS uvm_push_begin_fake(uvm_gpu_t *gpu, uvm_push_t *push);
 void uvm_push_end_fake(uvm_push_t *push);
 
-static bool uvm_push_is_fake(uvm_push_t *push)
-{
-    return !push->channel;
-}
-
 // Begin an inline data fragment in the push
 //
 // The inline data will be ignored by the GPU, but can be referenced from

kernel-open/nvidia-uvm/uvm_push_test.c

@@ -40,10 +40,9 @@
 
 static NvU32 get_push_begin_size(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_sec2(channel)) {
-        // SEC2 channels allocate CSL signature buffer at the beginning.
+    // SEC2 channels allocate CSL signature buffer at the beginning.
+    if (uvm_channel_is_sec2(channel))
         return UVM_CONF_COMPUTING_SIGN_BUF_MAX_SIZE + UVM_METHOD_SIZE;
-    }
 
     return 0;
 }
@@ -51,10 +50,14 @@ static NvU32 get_push_begin_size(uvm_channel_t *channel)
 // This is the storage required by a semaphore release.
 static NvU32 get_push_end_min_size(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_ce(channel)) {
-        if (uvm_channel_is_wlc(channel)) {
-            // Space (in bytes) used by uvm_push_end() on a Secure CE channel.
-            // Note that Secure CE semaphore release pushes two memset and one
+    uvm_gpu_t *gpu = uvm_channel_get_gpu(channel);
+
+    if (uvm_conf_computing_mode_enabled(gpu)) {
+        if (uvm_channel_is_ce(channel)) {
+            // Space (in bytes) used by uvm_push_end() on a CE channel when
+            // the Confidential Computing feature is enabled.
+            //
+            // Note that CE semaphore release pushes two memset and one
             // encryption method on top of the regular release.
             // Memset size
             // -------------
@@ -75,43 +78,44 @@ static NvU32 get_push_end_min_size(uvm_channel_t *channel)
             //
             // TOTAL                            : 144 Bytes
 
-            // Same as CE + LCIC GPPut update + LCIC doorbell
-            return 24 + 144 + 24 + 24;
-        }
-        else if (uvm_channel_is_secure_ce(channel)) {
+            if (uvm_channel_is_wlc(channel)) {
+                // Same as CE + LCIC GPPut update + LCIC doorbell
+                return 24 + 144 + 24 + 24;
+            }
+
             return 24 + 144;
         }
-        // Space (in bytes) used by uvm_push_end() on a CE channel.
-        return 24;
-    }
-    else if (uvm_channel_is_sec2(channel)) {
+
+        UVM_ASSERT(uvm_channel_is_sec2(channel));
+
         // A perfectly aligned inline buffer in SEC2 semaphore release.
         // We add UVM_METHOD_SIZE because of the NOP method to reserve
         // UVM_CSL_SIGN_AUTH_TAG_SIZE_BYTES (the inline buffer.)
         return 48 + UVM_CSL_SIGN_AUTH_TAG_SIZE_BYTES + UVM_METHOD_SIZE;
     }
 
-    return 0;
+    UVM_ASSERT(uvm_channel_is_ce(channel));
+
+    // Space (in bytes) used by uvm_push_end() on a CE channel.
+    return 24;
 }
 
 static NvU32 get_push_end_max_size(uvm_channel_t *channel)
 {
-    if (uvm_channel_is_ce(channel)) {
-        if (uvm_channel_is_wlc(channel)) {
-            // WLC pushes are always padded to UVM_MAX_WLC_PUSH_SIZE
-            return UVM_MAX_WLC_PUSH_SIZE;
-        }
-        // Space (in bytes) used by uvm_push_end() on a CE channel.
-        return get_push_end_min_size(channel);
-    }
-    else if (uvm_channel_is_sec2(channel)) {
-        // Space (in bytes) used by uvm_push_end() on a SEC2 channel.
-        // Note that SEC2 semaphore release uses an inline buffer with alignment
-        // requirements. This is the "worst" case semaphore_release storage.
-        return 48 + UVM_CSL_SIGN_AUTH_TAG_SIZE_BYTES + UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT;
-    }
+    // WLC pushes are always padded to UVM_MAX_WLC_PUSH_SIZE
+    if (uvm_channel_is_wlc(channel))
+        return UVM_MAX_WLC_PUSH_SIZE;
 
-    return 0;
+    // Space (in bytes) used by uvm_push_end() on a SEC2 channel.
+    // Note that SEC2 semaphore release uses an inline buffer with alignment
+    // requirements. This is the "worst" case semaphore_release storage.
+    if (uvm_channel_is_sec2(channel))
+        return 48 + UVM_CSL_SIGN_AUTH_TAG_SIZE_BYTES + UVM_CONF_COMPUTING_AUTH_TAG_ALIGNMENT;
+
+    UVM_ASSERT(uvm_channel_is_ce(channel));
+
+    // Space (in bytes) used by uvm_push_end() on a CE channel.
+    return get_push_end_min_size(channel);
 }
 
 static NV_STATUS test_push_end_size(uvm_va_space_t *va_space)
@@ -294,10 +298,19 @@ static NV_STATUS test_concurrent_pushes(uvm_va_space_t *va_space)
 {
     NV_STATUS status = NV_OK;
     uvm_gpu_t *gpu;
-    NvU32 i;
     uvm_push_t *pushes;
-    uvm_tracker_t tracker = UVM_TRACKER_INIT();
-    uvm_channel_type_t channel_type = UVM_CHANNEL_TYPE_GPU_INTERNAL;
+    uvm_tracker_t tracker;
+
+    // When the Confidential Computing feature is enabled, a channel reserved at
+    // the start of a push cannot be reserved again until that push ends. The
+    // test is waived, because the number of pushes it starts per pool exceeds
+    // the number of channels in the pool, so it would block indefinitely.
+    gpu = uvm_va_space_find_first_gpu(va_space);
+
+    if ((gpu != NULL) && uvm_conf_computing_mode_enabled(gpu))
+        return NV_OK;
+
+    uvm_tracker_init(&tracker);
 
     // As noted above, this test does unsafe things that would be detected by
     // lock tracking, opt-out.
@@ -310,16 +323,11 @@ static NV_STATUS test_concurrent_pushes(uvm_va_space_t *va_space)
     }
 
     for_each_va_space_gpu(gpu, va_space) {
+        NvU32 i;
 
-        // A secure channels reserved at the start of a push cannot be reserved
-        // again until that push ends. The test would block indefinitely
-        // if secure pools are not skipped, because the number of pushes started
-        // per pool exceeds the number of channels in the pool.
-        if (uvm_channel_type_requires_secure_pool(gpu, channel_type))
-            goto done;
         for (i = 0; i < UVM_PUSH_MAX_CONCURRENT_PUSHES; ++i) {
             uvm_push_t *push = &pushes[i];
-            status = uvm_push_begin(gpu->channel_manager, channel_type, push, "concurrent push %u", i);
+            status = uvm_push_begin(gpu->channel_manager, UVM_CHANNEL_TYPE_GPU_INTERNAL, push, "concurrent push %u", i);
             TEST_CHECK_GOTO(status == NV_OK, done);
         }
         for (i = 0; i < UVM_PUSH_MAX_CONCURRENT_PUSHES; ++i) {

kernel-open/nvidia-uvm/uvm_pushbuffer.c

@@ -458,7 +458,7 @@ static void decrypt_push(uvm_channel_t *channel, uvm_gpfifo_entry_t *gpfifo)
     void *push_unprotected_cpu_va;
     NvU32 pushbuffer_offset = gpfifo->pushbuffer_offset;
     NvU32 push_info_index = gpfifo->push_info - channel->push_infos;
-    uvm_pushbuffer_t *pushbuffer = channel->pool->manager->pushbuffer;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
     uvm_push_crypto_bundle_t *crypto_bundle = channel->conf_computing.push_crypto_bundles + push_info_index;
 
     if (channel->conf_computing.push_crypto_bundles == NULL)
@@ -499,7 +499,7 @@ void uvm_pushbuffer_mark_completed(uvm_channel_t *channel, uvm_gpfifo_entry_t *g
     uvm_pushbuffer_chunk_t *chunk;
     bool need_to_update_chunk = false;
     uvm_push_info_t *push_info = gpfifo->push_info;
-    uvm_pushbuffer_t *pushbuffer = channel->pool->manager->pushbuffer;
+    uvm_pushbuffer_t *pushbuffer = uvm_channel_get_pushbuffer(channel);
 
     UVM_ASSERT(gpfifo->type == UVM_GPFIFO_ENTRY_TYPE_NORMAL);
 

kernel-open/nvidia-uvm/uvm_range_group.c

@@ -264,7 +264,6 @@ NV_STATUS uvm_range_group_va_range_migrate(uvm_va_range_t *va_range,
         return NV_ERR_NO_MEMORY;
 
     uvm_assert_rwsem_locked(&va_range->va_space->lock);
-    va_block_context->policy = uvm_va_range_get_policy(va_range);
 
     // Iterate over blocks, populating them if necessary
     for (i = uvm_va_range_block_index(va_range, start); i <= uvm_va_range_block_index(va_range, end); ++i) {

kernel-open/nvidia-uvm/uvm_sec2_test.c

@@ -270,7 +270,7 @@ static NV_STATUS alloc_and_init_mem(uvm_gpu_t *gpu, uvm_mem_t **mem, size_t size
     *mem = NULL;
 
     if (type == MEM_ALLOC_TYPE_VIDMEM_PROTECTED) {
-        TEST_NV_CHECK_RET(uvm_mem_alloc_vidmem_protected(size, gpu, mem));
+        TEST_NV_CHECK_RET(uvm_mem_alloc_vidmem(size, gpu, mem));
         TEST_NV_CHECK_GOTO(uvm_mem_map_gpu_kernel(*mem, gpu), err);
         TEST_NV_CHECK_GOTO(ce_memset_gpu(gpu, *mem, size, 0xdead), err);
     }
@@ -348,9 +348,9 @@ static NV_STATUS cpu_decrypt(uvm_channel_t *channel,
     return NV_OK;
 }
 
-// gpu_encrypt uses a secure CE for encryption (instead of SEC2). SEC2 does not
-// support encryption. The following function is copied from uvm_ce_test.c and
-// adapted to SEC2 tests.
+// gpu_encrypt uses the Copy Engine for encryption, instead of SEC2. SEC2 does
+// not support encryption. The following function is copied from uvm_ce_test.c
+// and adapted to SEC2 tests.
 static void gpu_encrypt(uvm_push_t *push,
                         uvm_mem_t *dst_mem,
                         uvm_mem_t *src_mem,

kernel-open/nvidia-uvm/uvm_test.c

@@ -36,6 +36,7 @@
 #include "uvm_mmu.h"
 #include "uvm_gpu_access_counters.h"
 #include "uvm_pmm_sysmem.h"
+#include "uvm_migrate_pageable.h"
 
 static NV_STATUS uvm_test_get_gpu_ref_count(UVM_TEST_GET_GPU_REF_COUNT_PARAMS *params, struct file *filp)
 {
@@ -147,24 +148,23 @@ static NV_STATUS uvm_test_verify_bh_affinity(uvm_intr_handler_t *isr, int node)
 static NV_STATUS uvm_test_numa_check_affinity(UVM_TEST_NUMA_CHECK_AFFINITY_PARAMS *params, struct file *filp)
 {
     uvm_gpu_t *gpu;
-    NV_STATUS status;
-    uvm_rm_user_object_t user_rm_va_space = {
-        .rm_control_fd = -1,
-        .user_client = params->client,
-        .user_object = params->smc_part_ref
-    };
+    NV_STATUS status = NV_OK;
 
     if (!UVM_THREAD_AFFINITY_SUPPORTED())
         return NV_ERR_NOT_SUPPORTED;
 
-    status = uvm_gpu_retain_by_uuid(&params->gpu_uuid, &user_rm_va_space, &gpu);
-    if (status != NV_OK)
-        return status;
+    uvm_mutex_lock(&g_uvm_global.global_lock);
+
+    gpu = uvm_gpu_get_by_uuid(&params->gpu_uuid);
+    if (!gpu) {
+        status = NV_ERR_INVALID_DEVICE;
+        goto unlock;
+    }
 
     // If the GPU is not attached to a NUMA node, there is nothing to do.
     if (gpu->parent->closest_cpu_numa_node == NUMA_NO_NODE) {
         status = NV_ERR_NOT_SUPPORTED;
-        goto release;
+        goto unlock;
     }
 
     if (gpu->parent->replayable_faults_supported) {
@@ -173,7 +173,7 @@ static NV_STATUS uvm_test_numa_check_affinity(UVM_TEST_NUMA_CHECK_AFFINITY_PARAM
                                               gpu->parent->closest_cpu_numa_node);
         uvm_gpu_replayable_faults_isr_unlock(gpu->parent);
         if (status != NV_OK)
-            goto release;
+            goto unlock;
 
         if (gpu->parent->non_replayable_faults_supported) {
             uvm_gpu_non_replayable_faults_isr_lock(gpu->parent);
@@ -181,7 +181,7 @@ static NV_STATUS uvm_test_numa_check_affinity(UVM_TEST_NUMA_CHECK_AFFINITY_PARAM
                                                   gpu->parent->closest_cpu_numa_node);
             uvm_gpu_non_replayable_faults_isr_unlock(gpu->parent);
             if (status != NV_OK)
-                goto release;
+                goto unlock;
         }
 
         if (gpu->parent->access_counters_supported) {
@@ -191,8 +191,9 @@ static NV_STATUS uvm_test_numa_check_affinity(UVM_TEST_NUMA_CHECK_AFFINITY_PARAM
             uvm_gpu_access_counters_isr_unlock(gpu->parent);
         }
     }
-release:
-    uvm_gpu_release(gpu);
+
+unlock:
+    uvm_mutex_unlock(&g_uvm_global.global_lock);
     return status;
 }
 
@@ -331,6 +332,7 @@ long uvm_test_ioctl(struct file *filp, unsigned int cmd, unsigned long arg)
         UVM_ROUTE_CMD_STACK_NO_INIT_CHECK(UVM_TEST_CGROUP_ACCOUNTING_SUPPORTED, uvm_test_cgroup_accounting_supported);
         UVM_ROUTE_CMD_STACK_INIT_CHECK(UVM_TEST_SPLIT_INVALIDATE_DELAY, uvm_test_split_invalidate_delay);
         UVM_ROUTE_CMD_STACK_INIT_CHECK(UVM_TEST_CPU_CHUNK_API, uvm_test_cpu_chunk_api);
+        UVM_ROUTE_CMD_STACK_INIT_CHECK(UVM_TEST_SKIP_MIGRATE_VMA, uvm_test_skip_migrate_vma);
     }
 
     return -EINVAL;

kernel-open/nvidia-uvm/uvm_test_ioctl.h

@@ -28,6 +28,13 @@
 #include "uvm_ioctl.h"
 #include "nv_uvm_types.h"
 
+#define UVM_TEST_SKIP_MIGRATE_VMA                        UVM_TEST_IOCTL_BASE(103)
+typedef struct
+{
+    NvBool skip;                                         // In
+    NV_STATUS rmStatus;                                  // Out
+} UVM_TEST_SKIP_MIGRATE_VMA_PARAMS;
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -1196,8 +1203,6 @@ typedef struct
 typedef struct
 {
     NvProcessorUuid                 gpu_uuid;                                           // In
-    NvHandle                        client;                                             // In
-    NvHandle                        smc_part_ref;                                       // In
 
     NV_STATUS                       rmStatus;                                           // Out
 } UVM_TEST_NUMA_CHECK_AFFINITY_PARAMS;

kernel-open/nvidia-uvm/uvm_tools.c

@@ -229,6 +229,24 @@ static void unmap_user_pages(struct page **pages, void *addr, NvU64 size)
     uvm_kvfree(pages);
 }
 
+// This must be called with the mmap_lock held in read mode or better.
+static NV_STATUS check_vmas(struct mm_struct *mm, NvU64 start_va, NvU64 size)
+{
+    struct vm_area_struct *vma;
+    NvU64 addr = start_va;
+    NvU64 region_end = start_va + size;
+
+    do {
+        vma = find_vma(mm, addr);
+        if (!vma || !(addr >= vma->vm_start) || uvm_file_is_nvidia_uvm(vma->vm_file))
+            return NV_ERR_INVALID_ARGUMENT;
+
+        addr = vma->vm_end;
+    } while (addr < region_end);
+
+    return NV_OK;
+}
+
 // Map virtual memory of data from [user_va, user_va + size) of current process into kernel.
 // Sets *addr to kernel mapping and *pages to the array of struct pages that contain the memory.
 static NV_STATUS map_user_pages(NvU64 user_va, NvU64 size, void **addr, struct page ***pages)
@@ -237,7 +255,6 @@ static NV_STATUS map_user_pages(NvU64 user_va, NvU64 size, void **addr, struct p
     long ret = 0;
     long num_pages;
     long i;
-    struct vm_area_struct **vmas = NULL;
 
     *addr = NULL;
     *pages = NULL;
@@ -254,22 +271,30 @@ static NV_STATUS map_user_pages(NvU64 user_va, NvU64 size, void **addr, struct p
         goto fail;
     }
 
-    vmas = uvm_kvmalloc(sizeof(struct vm_area_struct *) * num_pages);
-    if (vmas == NULL) {
-        status = NV_ERR_NO_MEMORY;
+    // Although uvm_down_read_mmap_lock() is preferable due to its participation
+    // in the UVM lock dependency tracker, it cannot be used here. That's
+    // because pin_user_pages() may fault in HMM pages which are GPU-resident.
+    // When that happens, the UVM page fault handler would record another
+    // mmap_read_lock() on the same thread as this one, leading to a false
+    // positive lock dependency report.
+    //
+    // Therefore, use the lower level nv_mmap_read_lock() here.
+    nv_mmap_read_lock(current->mm);
+    status = check_vmas(current->mm, user_va, size);
+    if (status != NV_OK) {
+        nv_mmap_read_unlock(current->mm);
         goto fail;
     }
-
-    nv_mmap_read_lock(current->mm);
-    ret = NV_PIN_USER_PAGES(user_va, num_pages, FOLL_WRITE, *pages, vmas);
+    ret = NV_PIN_USER_PAGES(user_va, num_pages, FOLL_WRITE, *pages, NULL);
     nv_mmap_read_unlock(current->mm);
+
     if (ret != num_pages) {
         status = NV_ERR_INVALID_ARGUMENT;
         goto fail;
     }
 
     for (i = 0; i < num_pages; i++) {
-        if (page_count((*pages)[i]) > MAX_PAGE_COUNT || uvm_file_is_nvidia_uvm(vmas[i]->vm_file)) {
+        if (page_count((*pages)[i]) > MAX_PAGE_COUNT) {
             status = NV_ERR_INVALID_ARGUMENT;
             goto fail;
         }
@@ -279,15 +304,12 @@ static NV_STATUS map_user_pages(NvU64 user_va, NvU64 size, void **addr, struct p
     if (*addr == NULL)
         goto fail;
 
-    uvm_kvfree(vmas);
     return NV_OK;
 
 fail:
     if (*pages == NULL)
         return status;
 
-    uvm_kvfree(vmas);
-
     if (ret > 0)
         uvm_put_user_pages_dirty(*pages, ret);
     else if (ret < 0)
@@ -1060,25 +1082,19 @@ void uvm_tools_broadcast_replay(uvm_gpu_t *gpu,
 }
 
 
-void uvm_tools_broadcast_replay_sync(uvm_gpu_t *gpu,
-                                     NvU32 batch_id,
-                                     uvm_fault_client_type_t client_type)
+void uvm_tools_broadcast_replay_sync(uvm_gpu_t *gpu, NvU32 batch_id, uvm_fault_client_type_t client_type)
 {
     UVM_ASSERT(!gpu->parent->has_clear_faulted_channel_method);
 
     if (!tools_is_event_enabled_in_any_va_space(UvmEventTypeGpuFaultReplay))
         return;
 
-    record_replay_event_helper(gpu->id,
-                               batch_id,
-                               client_type,
-                               NV_GETTIME(),
-                               gpu->parent->host_hal->get_time(gpu));
+    record_replay_event_helper(gpu->id, batch_id, client_type, NV_GETTIME(), gpu->parent->host_hal->get_time(gpu));
 }
 
 void uvm_tools_broadcast_access_counter(uvm_gpu_t *gpu,
                                         const uvm_access_counter_buffer_entry_t *buffer_entry,
-                                        bool on_managed)
+                                        bool on_managed_phys)
 {
     UvmEventEntry entry;
     UvmEventTestAccessCounterInfo *info = &entry.testEventData.accessCounter;
@@ -1097,6 +1113,7 @@ void uvm_tools_broadcast_access_counter(uvm_gpu_t *gpu,
     info->srcIndex            = uvm_id_value(gpu->id);
     info->address             = buffer_entry->address.address;
     info->isVirtual           = buffer_entry->address.is_virtual? 1: 0;
+
     if (buffer_entry->address.is_virtual) {
         info->instancePtr         = buffer_entry->virtual_info.instance_ptr.address;
         info->instancePtrAperture = g_hal_to_tools_aperture_table[buffer_entry->virtual_info.instance_ptr.aperture];
@@ -1104,9 +1121,10 @@ void uvm_tools_broadcast_access_counter(uvm_gpu_t *gpu,
     }
     else {
         info->aperture            = g_hal_to_tools_aperture_table[buffer_entry->address.aperture];
+        info->physOnManaged       = on_managed_phys? 1 : 0;
     }
+
     info->isFromCpu           = buffer_entry->counter_type == UVM_ACCESS_COUNTER_TYPE_MOMC? 1: 0;
-    info->onManaged           = on_managed? 1 : 0;
     info->value               = buffer_entry->counter_value;
     info->subGranularity      = buffer_entry->sub_granularity;
     info->bank                = buffer_entry->bank;
@@ -2047,7 +2065,11 @@ static NV_STATUS tools_access_process_memory(uvm_va_space_t *va_space,
 
         // The RM flavor of the lock is needed to perform ECC checks.
         uvm_va_space_down_read_rm(va_space);
-        status = uvm_va_block_find_create(va_space, UVM_PAGE_ALIGN_DOWN(target_va_start), block_context, &block);
+        if (mm)
+            status = uvm_va_block_find_create(va_space, UVM_PAGE_ALIGN_DOWN(target_va_start), &block_context->hmm.vma, &block);
+        else
+            status = uvm_va_block_find_create_managed(va_space, UVM_PAGE_ALIGN_DOWN(target_va_start), &block);
+
         if (status != NV_OK)
             goto unlock_and_exit;
 

kernel-open/nvidia-uvm/uvm_tools.h

@@ -102,18 +102,13 @@ void uvm_tools_record_read_duplicate_invalidate(uvm_va_block_t *va_block,
                                                 uvm_va_block_region_t region,
                                                 const uvm_page_mask_t *page_mask);
 
-void uvm_tools_broadcast_replay(uvm_gpu_t *gpu,
-                                uvm_push_t *push,
-                                NvU32 batch_id,
-                                uvm_fault_client_type_t client_type);
+void uvm_tools_broadcast_replay(uvm_gpu_t *gpu, uvm_push_t *push, NvU32 batch_id, uvm_fault_client_type_t client_type);
 
-void uvm_tools_broadcast_replay_sync(uvm_gpu_t *gpu,
-                                     NvU32 batch_id,
-                                     uvm_fault_client_type_t client_type);
+void uvm_tools_broadcast_replay_sync(uvm_gpu_t *gpu, NvU32 batch_id, uvm_fault_client_type_t client_type);
 
 void uvm_tools_broadcast_access_counter(uvm_gpu_t *gpu,
                                         const uvm_access_counter_buffer_entry_t *buffer_entry,
-                                        bool on_managed);
+                                        bool on_managed_phys);
 
 void uvm_tools_test_hmm_split_invalidate(uvm_va_space_t *va_space);
 

kernel-open/nvidia-uvm/uvm_turing.c

@@ -1,5 +1,5 @@
 /*******************************************************************************
-    Copyright (c) 2017-2021 NVIDIA Corporation
+    Copyright (c) 2017-2023 NVIDIA Corporation
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to
@@ -93,4 +93,6 @@ void uvm_hal_turing_arch_init_properties(uvm_parent_gpu_t *parent_gpu)
     parent_gpu->smc.supported = false;
 
     parent_gpu->plc_supported = true;
+
+    parent_gpu->no_ats_range_required = false;
 }

kernel-open/nvidia-uvm/uvm_types.h

@@ -967,8 +967,10 @@ typedef struct
     NvU8 isFromCpu;
 
     NvU8 veId;
-    NvU8 onManaged;         // The access counter notification was triggered on
-                            // a managed memory region
+
+    // The physical access counter notification was triggered on a managed
+    // memory region. This is not set for virtual access counter notifications.
+    NvU8 physOnManaged;
 
     NvU32 value;
     NvU32 subGranularity;

kernel-open/nvidia-uvm/uvm_va_block.c

@@ -106,36 +106,6 @@ uvm_va_space_t *uvm_va_block_get_va_space(uvm_va_block_t *va_block)
     return va_space;
 }
 
-bool uvm_va_block_check_policy_is_valid(uvm_va_block_t *va_block,
-                                        const uvm_va_policy_t *policy,
-                                        uvm_va_block_region_t region)
-{
-    uvm_assert_mutex_locked(&va_block->lock);
-
-    if (uvm_va_block_is_hmm(va_block)) {
-        const uvm_va_policy_node_t *node;
-
-        if (uvm_va_policy_is_default(policy)) {
-            // There should only be the default policy within the region.
-            node = uvm_va_policy_node_iter_first(va_block,
-                                                 uvm_va_block_region_start(va_block, region),
-                                                 uvm_va_block_region_end(va_block, region));
-            UVM_ASSERT(!node);
-        }
-        else {
-            // The policy node should cover the region.
-            node = uvm_va_policy_node_from_policy(policy);
-            UVM_ASSERT(node->node.start <= uvm_va_block_region_start(va_block, region));
-            UVM_ASSERT(node->node.end >= uvm_va_block_region_end(va_block, region));
-        }
-    }
-    else {
-        UVM_ASSERT(policy == uvm_va_range_get_policy(va_block->va_range));
-    }
-
-    return true;
-}
-
 static NvU64 block_gpu_pte_flag_cacheable(uvm_va_block_t *block, uvm_gpu_t *gpu, uvm_processor_id_t resident_id)
 {
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(block);
@@ -1790,6 +1760,21 @@ static NvU32 block_phys_page_size(uvm_va_block_t *block, block_phys_page_t page)
     return (NvU32)chunk_size;
 }
 
+NvU32 uvm_va_block_get_physical_size(uvm_va_block_t *block,
+                                     uvm_processor_id_t processor,
+                                     uvm_page_index_t page_index)
+{
+    block_phys_page_t page;
+
+    UVM_ASSERT(block);
+
+    uvm_assert_mutex_locked(&block->lock);
+
+    page = block_phys_page(processor, page_index);
+
+    return block_phys_page_size(block, page);
+}
+
 static uvm_pte_bits_cpu_t get_cpu_pte_bit_index(uvm_prot_t prot)
 {
     uvm_pte_bits_cpu_t pte_bit_index = UVM_PTE_BITS_CPU_MAX;
@@ -3055,7 +3040,7 @@ static NV_STATUS conf_computing_copy_pages_finish(uvm_va_block_t *block,
     void *auth_tag_buffer_base = uvm_mem_get_cpu_addr_kernel(dma_buffer->auth_tag);
     void *staging_buffer_base = uvm_mem_get_cpu_addr_kernel(dma_buffer->alloc);
 
-    UVM_ASSERT(uvm_channel_is_secure(push->channel));
+    UVM_ASSERT(uvm_conf_computing_mode_enabled(push->gpu));
 
     if (UVM_ID_IS_GPU(copy_state->dst.id))
         return NV_OK;
@@ -3106,7 +3091,7 @@ static void block_copy_push(uvm_va_block_t *block,
 
     uvm_push_set_flag(push, UVM_PUSH_FLAG_NEXT_MEMBAR_NONE);
 
-    if (uvm_channel_is_secure(push->channel)) {
+    if (uvm_conf_computing_mode_enabled(gpu)) {
         if (UVM_ID_IS_CPU(copy_state->src.id))
             conf_computing_block_copy_push_cpu_to_gpu(block, copy_state, region, push);
         else
@@ -3134,19 +3119,18 @@ static NV_STATUS block_copy_end_push(uvm_va_block_t *block,
     //       at that point.
     uvm_push_end(push);
 
-    if ((push_status == NV_OK) && uvm_channel_is_secure(push->channel))
+    if ((push_status == NV_OK) && uvm_conf_computing_mode_enabled(push->gpu))
         push_status = conf_computing_copy_pages_finish(block, copy_state, push);
 
     tracker_status = uvm_tracker_add_push_safe(copy_tracker, push);
     if (push_status == NV_OK)
         push_status = tracker_status;
 
-    if (uvm_channel_is_secure(push->channel)) {
-        uvm_gpu_t *gpu = uvm_push_get_gpu(push);
+    if (uvm_conf_computing_mode_enabled(push->gpu)) {
         uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
 
         uvm_tracker_overwrite_with_push(&local_tracker, push);
-        uvm_conf_computing_dma_buffer_free(&gpu->conf_computing.dma_buffer_pool,
+        uvm_conf_computing_dma_buffer_free(&push->gpu->conf_computing.dma_buffer_pool,
                                            copy_state->dma_buffer,
                                            &local_tracker);
         copy_state->dma_buffer = NULL;
@@ -3698,7 +3682,6 @@ NV_STATUS uvm_va_block_make_resident_copy(uvm_va_block_t *va_block,
 
     uvm_assert_mutex_locked(&va_block->lock);
     UVM_ASSERT(uvm_va_block_is_hmm(va_block) || va_block->va_range->type == UVM_VA_RANGE_TYPE_MANAGED);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, va_block_context->policy, region));
 
     resident_mask = block_resident_mask_get_alloc(va_block, dest_id);
     if (!resident_mask)
@@ -3945,7 +3928,6 @@ NV_STATUS uvm_va_block_make_resident_read_duplicate(uvm_va_block_t *va_block,
 
     // TODO: Bug 3660922: need to implement HMM read duplication support.
     UVM_ASSERT(!uvm_va_block_is_hmm(va_block));
-    UVM_ASSERT(va_block_context->policy == uvm_va_range_get_policy(va_block->va_range));
 
     va_block_context->make_resident.dest_id = dest_id;
     va_block_context->make_resident.cause = cause;
@@ -4743,7 +4725,7 @@ static void block_unmap_cpu(uvm_va_block_t *block, uvm_va_block_region_t region,
 // Given a mask of mapped pages, returns true if any of the pages in the mask
 // are mapped remotely by the given GPU.
 static bool block_has_remote_mapping_gpu(uvm_va_block_t *block,
-                                         uvm_va_block_context_t *block_context,
+                                         uvm_page_mask_t *scratch_page_mask,
                                          uvm_gpu_id_t gpu_id,
                                          const uvm_page_mask_t *mapped_pages)
 {
@@ -4765,7 +4747,7 @@ static bool block_has_remote_mapping_gpu(uvm_va_block_t *block,
     }
 
     // Remote pages are pages which are mapped but not resident locally
-    return uvm_page_mask_andnot(&block_context->scratch_page_mask, mapped_pages, &gpu_state->resident);
+    return uvm_page_mask_andnot(scratch_page_mask, mapped_pages, &gpu_state->resident);
 }
 
 // Writes pte_clear_val to the 4k PTEs covered by clear_page_mask. If
@@ -6660,7 +6642,7 @@ static NV_STATUS block_unmap_gpu(uvm_va_block_t *block,
     if (status != NV_OK)
         return status;
 
-    only_local_mappings = !block_has_remote_mapping_gpu(block, block_context, gpu->id, pages_to_unmap);
+    only_local_mappings = !block_has_remote_mapping_gpu(block, &block_context->scratch_page_mask, gpu->id, pages_to_unmap);
     tlb_membar = uvm_hal_downgrade_membar_type(gpu, only_local_mappings);
 
     status = uvm_push_begin_acquire(gpu->channel_manager,
@@ -6795,16 +6777,15 @@ static NV_STATUS uvm_cpu_insert_page(struct vm_area_struct *vma,
 }
 
 static uvm_prot_t compute_logical_prot(uvm_va_block_t *va_block,
-                                       uvm_va_block_context_t *va_block_context,
+                                       struct vm_area_struct *hmm_vma,
                                        uvm_page_index_t page_index)
 {
-    struct vm_area_struct *vma;
     uvm_prot_t logical_prot;
 
     if (uvm_va_block_is_hmm(va_block)) {
         NvU64 addr = uvm_va_block_cpu_page_address(va_block, page_index);
 
-        logical_prot = uvm_hmm_compute_logical_prot(va_block, va_block_context, addr);
+        logical_prot = uvm_hmm_compute_logical_prot(va_block, hmm_vma, addr);
     }
     else {
         uvm_va_range_t *va_range = va_block->va_range;
@@ -6816,6 +6797,8 @@ static uvm_prot_t compute_logical_prot(uvm_va_block_t *va_block,
             logical_prot = UVM_PROT_NONE;
         }
         else {
+            struct vm_area_struct *vma;
+
             vma = uvm_va_range_vma(va_range);
 
             if (!(vma->vm_flags & VM_READ))
@@ -6865,13 +6848,15 @@ static struct page *block_page_get(uvm_va_block_t *block, block_phys_page_t bloc
 //    with new_prot permissions
 //  - Guarantee that vm_insert_page is safe to use (vma->vm_mm has a reference
 //    and mmap_lock is held in at least read mode)
+//  - For HMM blocks that vma is valid and safe to use, vma->vm_mm has a
+//    reference and mmap_lock is held in at least read mode
 //  - Ensure that the struct page corresponding to the physical memory being
 //    mapped exists
 //  - Manage the block's residency bitmap
 //  - Ensure that the block hasn't been killed (block->va_range is present)
 //  - Update the pte/mapping tracking state on success
 static NV_STATUS block_map_cpu_page_to(uvm_va_block_t *block,
-                                       uvm_va_block_context_t *va_block_context,
+                                       struct vm_area_struct *hmm_vma,
                                        uvm_processor_id_t resident_id,
                                        uvm_page_index_t page_index,
                                        uvm_prot_t new_prot)
@@ -6884,7 +6869,7 @@ static NV_STATUS block_map_cpu_page_to(uvm_va_block_t *block,
     NvU64 addr;
     struct page *page;
 
-    UVM_ASSERT(uvm_va_block_is_hmm(block) || va_range->type == UVM_VA_RANGE_TYPE_MANAGED);
+    UVM_ASSERT((uvm_va_block_is_hmm(block) && hmm_vma) || va_range->type == UVM_VA_RANGE_TYPE_MANAGED);
     UVM_ASSERT(new_prot != UVM_PROT_NONE);
     UVM_ASSERT(new_prot < UVM_PROT_MAX);
     UVM_ASSERT(uvm_processor_mask_test(&va_space->accessible_from[uvm_id_value(resident_id)], UVM_ID_CPU));
@@ -6905,7 +6890,7 @@ static NV_STATUS block_map_cpu_page_to(uvm_va_block_t *block,
 
     // Check for existing VMA permissions. They could have been modified after
     // the initial mmap by mprotect.
-    if (new_prot > compute_logical_prot(block, va_block_context, page_index))
+    if (new_prot > compute_logical_prot(block, hmm_vma, page_index))
         return NV_ERR_INVALID_ACCESS_TYPE;
 
     if (uvm_va_block_is_hmm(block)) {
@@ -7002,7 +6987,7 @@ static NV_STATUS block_map_cpu_to(uvm_va_block_t *block,
 
     for_each_va_block_page_in_region_mask(page_index, pages_to_map, region) {
         status = block_map_cpu_page_to(block,
-                                       block_context,
+                                       block_context->hmm.vma,
                                        resident_id,
                                        page_index,
                                        new_prot);
@@ -7235,13 +7220,13 @@ NV_STATUS uvm_va_block_map(uvm_va_block_t *va_block,
     const uvm_page_mask_t *pte_mask;
     uvm_page_mask_t *running_page_mask = &va_block_context->mapping.map_running_page_mask;
     NV_STATUS status;
+    const uvm_va_policy_t *policy = uvm_va_policy_get_region(va_block, region);
 
     va_block_context->mapping.cause = cause;
 
     UVM_ASSERT(new_prot != UVM_PROT_NONE);
     UVM_ASSERT(new_prot < UVM_PROT_MAX);
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, va_block_context->policy, region));
 
     // Mapping is not supported on the eviction path that doesn't hold the VA
     // space lock.
@@ -7283,7 +7268,7 @@ NV_STATUS uvm_va_block_map(uvm_va_block_t *va_block,
 
     // Map per resident location so we can more easily detect physically-
     // contiguous mappings.
-    map_get_allowed_destinations(va_block, va_block_context, va_block_context->policy, id, &allowed_destinations);
+    map_get_allowed_destinations(va_block, va_block_context, policy, id, &allowed_destinations);
 
     for_each_closest_id(resident_id, &allowed_destinations, id, va_space) {
         if (UVM_ID_IS_CPU(id)) {
@@ -7589,8 +7574,6 @@ NV_STATUS uvm_va_block_map_mask(uvm_va_block_t *va_block,
     NV_STATUS tracker_status;
     uvm_processor_id_t id;
 
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, va_block_context->policy, region));
-
     for_each_id_in_mask(id, map_processor_mask) {
         status = uvm_va_block_map(va_block,
                                   va_block_context,
@@ -8280,14 +8263,6 @@ void uvm_va_block_munmap_region(uvm_va_block_t *va_block,
     event_data.block_munmap.region = region;
     uvm_perf_event_notify(&va_space->perf_events, UVM_PERF_EVENT_BLOCK_MUNMAP, &event_data);
 
-    // Set a flag so that GPU fault events are flushed since they might refer
-    // to the region being unmapped.
-    // Note that holding the va_block lock prevents GPU VA spaces from
-    // being removed so the registered_gpu_va_spaces mask is stable.
-    for_each_gpu_id_in_mask(gpu_id, &va_space->registered_gpu_va_spaces) {
-        uvm_processor_mask_set_atomic(&va_space->needs_fault_buffer_flush, gpu_id);
-    }
-
     // Release any remaining vidmem chunks in the given region.
     for_each_gpu_id(gpu_id) {
         uvm_va_block_gpu_state_t *gpu_state = uvm_va_block_gpu_state_get(va_block, gpu_id);
@@ -9574,7 +9549,7 @@ static bool block_region_might_read_duplicate(uvm_va_block_t *va_block,
 //       could be changed in the future to optimize multiple faults/counters on
 //       contiguous pages.
 static uvm_prot_t compute_new_permission(uvm_va_block_t *va_block,
-                                         uvm_va_block_context_t *va_block_context,
+                                         struct vm_area_struct *hmm_vma,
                                          uvm_page_index_t page_index,
                                          uvm_processor_id_t fault_processor_id,
                                          uvm_processor_id_t new_residency,
@@ -9587,7 +9562,7 @@ static uvm_prot_t compute_new_permission(uvm_va_block_t *va_block,
     //       query_promote: upgrade access privileges to avoid future faults IF
     //       they don't trigger further revocations.
     new_prot = uvm_fault_access_type_to_prot(access_type);
-    logical_prot = compute_logical_prot(va_block, va_block_context, page_index);
+    logical_prot = compute_logical_prot(va_block, hmm_vma, page_index);
 
     UVM_ASSERT(logical_prot >= new_prot);
 
@@ -9612,15 +9587,9 @@ static uvm_prot_t compute_new_permission(uvm_va_block_t *va_block,
         if (uvm_processor_mask_empty(&revoke_processors))
             new_prot = UVM_PROT_READ_WRITE;
     }
-    if (logical_prot == UVM_PROT_READ_WRITE_ATOMIC) {
-        // HMM allocations with logical read/write/atomic permission can be
-        // upgraded without notifying the driver so assume read/write/atomic
-        // even if the fault is only for reading.
-        if (new_prot == UVM_PROT_READ_WRITE ||
-            (UVM_ID_IS_CPU(fault_processor_id) && uvm_va_block_is_hmm(va_block))) {
-            if (uvm_processor_mask_test(&va_space->has_native_atomics[uvm_id_value(new_residency)], fault_processor_id))
-                new_prot = UVM_PROT_READ_WRITE_ATOMIC;
-        }
+    if (logical_prot == UVM_PROT_READ_WRITE_ATOMIC && new_prot == UVM_PROT_READ_WRITE) {
+        if (uvm_processor_mask_test(&va_space->has_native_atomics[uvm_id_value(new_residency)], fault_processor_id))
+            new_prot = UVM_PROT_READ_WRITE_ATOMIC;
     }
 
     return new_prot;
@@ -9736,11 +9705,10 @@ NV_STATUS uvm_va_block_add_mappings_after_migration(uvm_va_block_t *va_block,
     uvm_va_space_t *va_space = uvm_va_block_get_va_space(va_block);
     const uvm_page_mask_t *final_page_mask = map_page_mask;
     uvm_tracker_t local_tracker = UVM_TRACKER_INIT();
-    const uvm_va_policy_t *policy = va_block_context->policy;
+    const uvm_va_policy_t *policy = uvm_va_policy_get_region(va_block, region);
     uvm_processor_id_t preferred_location;
 
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, policy, region));
 
     // Read duplication takes precedence over SetAccessedBy.
     //
@@ -9857,8 +9825,6 @@ out:
     return status == NV_OK ? tracker_status : status;
 }
 
-// TODO: Bug 1750144: check logical permissions from HMM to know what's the
-//       maximum allowed.
 uvm_prot_t uvm_va_block_page_compute_highest_permission(uvm_va_block_t *va_block,
                                                         uvm_processor_id_t processor_id,
                                                         uvm_page_index_t page_index)
@@ -9935,14 +9901,18 @@ uvm_prot_t uvm_va_block_page_compute_highest_permission(uvm_va_block_t *va_block
         // Exclude the processor for which the mapping protections are being computed
         uvm_processor_mask_clear(&write_mappings, processor_id);
 
-        // At this point, any processor with atomic mappings either has native atomics support to the
-        // processor with the resident copy or has disabled system-wide atomics. If the requesting
-        // processor has disabled system-wide atomics or has native atomics to that processor, we can
-        // map with ATOMIC privileges. Likewise, if there are no other processors with WRITE or ATOMIC
-        // mappings, we can map with ATOMIC privileges.
+        // At this point, any processor with atomic mappings either has native
+        // atomics support to the processor with the resident copy or has
+        // disabled system-wide atomics. If the requesting processor has
+        // disabled system-wide atomics or has native atomics to that processor,
+        // we can map with ATOMIC privileges. Likewise, if there are no other
+        // processors with WRITE or ATOMIC mappings, we can map with ATOMIC
+        // privileges. For HMM, don't allow GPU atomic access to remote mapped
+        // system memory even if there are no write mappings since CPU access
+        // can be upgraded without notification.
         if (!uvm_processor_mask_test(&va_space->system_wide_atomics_enabled_processors, processor_id) ||
             uvm_processor_mask_test(&va_space->has_native_atomics[uvm_id_value(residency)], processor_id) ||
-            uvm_processor_mask_empty(&write_mappings)) {
+            (uvm_processor_mask_empty(&write_mappings) && !uvm_va_block_is_hmm(va_block))) {
             return UVM_PROT_READ_WRITE_ATOMIC;
         }
 
@@ -9964,8 +9934,6 @@ NV_STATUS uvm_va_block_add_mappings(uvm_va_block_t *va_block,
     uvm_range_group_range_iter_t iter;
     uvm_prot_t prot_to_map;
 
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block, va_block_context->policy, region));
-
     if (UVM_ID_IS_CPU(processor_id) && !uvm_va_block_is_hmm(va_block)) {
         if (!uvm_va_range_vma_check(va_range, va_block_context->mm))
             return NV_OK;
@@ -10194,6 +10162,34 @@ static uvm_processor_id_t block_select_residency(uvm_va_block_t *va_block,
         uvm_processor_mask_test(&va_space->accessible_from[uvm_id_value(preferred_location)], processor_id))
         return preferred_location;
 
+    // Check if we should map the closest resident processor remotely on remote CPU fault
+    //
+    // When faulting on CPU, there's a linux process on behalf of it, which is associated
+    // with a unique VM pointed by current->mm. A block of memory residing on GPU is also
+    // associated with VM, pointed by va_block_context->mm. If they match, it's a regular
+    // (local) fault, and we may want to migrate a page from GPU to CPU.
+    // If it's a 'remote' fault, i.e. linux process differs from one associated with block
+    // VM, we might preserve residence.
+    //
+    // Establishing a remote fault without access counters means the memory could stay in
+    // the wrong spot for a long time, which is why we prefer to avoid creating remote
+    // mappings. However when NIC accesses a memory residing on GPU, it's worth to keep it
+    // in place for NIC accesses.
+    //
+    // The logic that's used to detect remote faulting also keeps memory in place for
+    // ptrace accesses. We would prefer to control those policies separately, but the
+    // NIC case takes priority.
+    // If the accessing processor is CPU, we're either handling a fault
+    // from other than owning process, or we're handling an MOMC
+    // notification. Only prevent migration for the former.
+    if (UVM_ID_IS_CPU(processor_id) &&
+        operation != UVM_SERVICE_OPERATION_ACCESS_COUNTERS &&        
+        uvm_processor_mask_test(&va_space->accessible_from[uvm_id_value(closest_resident_processor)], processor_id) &&
+        va_block_context->mm != current->mm) {
+        UVM_ASSERT(va_block_context->mm != NULL);
+        return closest_resident_processor;
+    }
+
     // If the page is resident on a processor other than the preferred location,
     // or the faulting processor can't access the preferred location, we select
     // the faulting processor as the new residency.
@@ -10212,11 +10208,8 @@ uvm_processor_id_t uvm_va_block_select_residency(uvm_va_block_t *va_block,
 {
     uvm_processor_id_t id;
 
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block,
-                                                  va_block_context->policy,
-                                                  uvm_va_block_region_for_page(page_index)));
     UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block,
-                                                  va_block_context,
+                                                  va_block_context->hmm.vma,
                                                   uvm_va_block_region_for_page(page_index)));
 
     id = block_select_residency(va_block,
@@ -10260,6 +10253,7 @@ static bool check_access_counters_dont_revoke(uvm_va_block_t *block,
 // Update service_context->prefetch_hint, service_context->per_processor_masks,
 // and service_context->region.
 static void uvm_va_block_get_prefetch_hint(uvm_va_block_t *va_block,
+                                           const uvm_va_policy_t *policy,
                                            uvm_service_block_context_t *service_context)
 {
     uvm_processor_id_t new_residency;
@@ -10270,20 +10264,19 @@ static void uvm_va_block_get_prefetch_hint(uvm_va_block_t *va_block,
     if (uvm_processor_mask_get_count(&service_context->resident_processors) == 1) {
         uvm_page_index_t page_index;
         uvm_page_mask_t *new_residency_mask;
-        const uvm_va_policy_t *policy = service_context->block_context.policy;
 
         new_residency = uvm_processor_mask_find_first_id(&service_context->resident_processors);
         new_residency_mask = &service_context->per_processor_masks[uvm_id_value(new_residency)].new_residency;
 
         // Update prefetch tracking structure with the pages that will migrate
         // due to faults
-        uvm_perf_prefetch_get_hint(va_block,
-                                   &service_context->block_context,
-                                   new_residency,
-                                   new_residency_mask,
-                                   service_context->region,
-                                   &service_context->prefetch_bitmap_tree,
-                                   &service_context->prefetch_hint);
+        uvm_perf_prefetch_get_hint_va_block(va_block,
+                                            &service_context->block_context,
+                                            new_residency,
+                                            new_residency_mask,
+                                            service_context->region,
+                                            &service_context->prefetch_bitmap_tree,
+                                            &service_context->prefetch_hint);
 
         // Obtain the prefetch hint and give a fake fault access type to the
         // prefetched pages
@@ -10468,7 +10461,7 @@ NV_STATUS uvm_va_block_service_finish(uvm_processor_id_t processor_id,
 
     for_each_va_block_page_in_region_mask(page_index, new_residency_mask, service_context->region) {
         new_prot = compute_new_permission(va_block,
-                                          &service_context->block_context,
+                                          service_context->block_context.hmm.vma,
                                           page_index,
                                           processor_id,
                                           new_residency,
@@ -10711,11 +10704,8 @@ NV_STATUS uvm_va_block_service_locked(uvm_processor_id_t processor_id,
     NV_STATUS status = NV_OK;
 
     uvm_assert_mutex_locked(&va_block->lock);
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block,
-                                                  service_context->block_context.policy,
-                                                  service_context->region));
     UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block,
-                                                  &service_context->block_context,
+                                                  service_context->block_context.hmm.vma,
                                                   service_context->region));
 
     // GPU fault servicing must be done under the VA space read lock. GPU fault
@@ -10729,7 +10719,9 @@ NV_STATUS uvm_va_block_service_locked(uvm_processor_id_t processor_id,
     else
         uvm_assert_rwsem_locked_read(&va_space->lock);
 
-    uvm_va_block_get_prefetch_hint(va_block, service_context);
+    uvm_va_block_get_prefetch_hint(va_block,
+                                   uvm_va_policy_get_region(va_block, service_context->region),
+                                   service_context);
 
     for_each_id_in_mask(new_residency, &service_context->resident_processors) {
         if (uvm_va_block_is_hmm(va_block)) {
@@ -10756,17 +10748,14 @@ NV_STATUS uvm_va_block_check_logical_permissions(uvm_va_block_t *va_block,
                                                  uvm_va_block_context_t *va_block_context,
                                                  uvm_processor_id_t processor_id,
                                                  uvm_page_index_t page_index,
-                                                 uvm_fault_type_t access_type,
+                                                 uvm_fault_access_type_t access_type,
                                                  bool allow_migration)
 {
     uvm_va_range_t *va_range = va_block->va_range;
     uvm_prot_t access_prot = uvm_fault_access_type_to_prot(access_type);
 
-    UVM_ASSERT(uvm_va_block_check_policy_is_valid(va_block,
-                                                  va_block_context->policy,
-                                                  uvm_va_block_region_for_page(page_index)));
     UVM_ASSERT(uvm_hmm_check_context_vma_is_valid(va_block,
-                                                  va_block_context,
+                                                  va_block_context->hmm.vma,
                                                   uvm_va_block_region_for_page(page_index)));
 
     // CPU permissions are checked later by block_map_cpu_page.
@@ -10784,8 +10773,8 @@ NV_STATUS uvm_va_block_check_logical_permissions(uvm_va_block_t *va_block,
         // vm_flags at any moment (for example on mprotect) and here we are not
         // guaranteed to have vma->vm_mm->mmap_lock. During tests we ensure that
         // this scenario does not happen.
-        if ((va_block_context->mm || uvm_enable_builtin_tests) &&
-            (access_prot > compute_logical_prot(va_block, va_block_context, page_index)))
+        if (((va_block->hmm.va_space && va_block->hmm.va_space->va_space_mm.mm) || uvm_enable_builtin_tests) &&
+            (access_prot > compute_logical_prot(va_block, va_block_context->hmm.vma, page_index)))
             return NV_ERR_INVALID_ACCESS_TYPE;
     }
 
@@ -10871,6 +10860,7 @@ static NV_STATUS block_cpu_fault_locked(uvm_va_block_t *va_block,
     uvm_perf_thrashing_hint_t thrashing_hint;
     uvm_processor_id_t new_residency;
     bool read_duplicate;
+    const uvm_va_policy_t *policy;
 
     uvm_assert_rwsem_locked(&va_space->lock);
 
@@ -10879,13 +10869,13 @@ static NV_STATUS block_cpu_fault_locked(uvm_va_block_t *va_block,
 
     uvm_assert_mmap_lock_locked(service_context->block_context.mm);
 
-    service_context->block_context.policy = uvm_va_policy_get(va_block, fault_addr);
+    policy = uvm_va_policy_get(va_block, fault_addr);
 
     if (service_context->num_retries == 0) {
         // notify event to tools/performance heuristics
         uvm_perf_event_notify_cpu_fault(&va_space->perf_events,
                                         va_block,
-                                        service_context->block_context.policy->preferred_location,
+                                        policy->preferred_location,
                                         fault_addr,
                                         fault_access_type > UVM_FAULT_ACCESS_TYPE_READ,
                                         KSTK_EIP(current));
@@ -10930,7 +10920,7 @@ static NV_STATUS block_cpu_fault_locked(uvm_va_block_t *va_block,
                                                   page_index,
                                                   UVM_ID_CPU,
                                                   uvm_fault_access_type_mask_bit(fault_access_type),
-                                                  service_context->block_context.policy,
+                                                  policy,
                                                   &thrashing_hint,
                                                   UVM_SERVICE_OPERATION_REPLAYABLE_FAULTS,
                                                   &read_duplicate);
@@ -11030,7 +11020,6 @@ NV_STATUS uvm_va_block_find(uvm_va_space_t *va_space, NvU64 addr, uvm_va_block_t
 NV_STATUS uvm_va_block_find_create_in_range(uvm_va_space_t *va_space,
                                             uvm_va_range_t *va_range,
                                             NvU64 addr,
-                                            uvm_va_block_context_t *va_block_context,
                                             uvm_va_block_t **out_block)
 {
     size_t index;
@@ -11038,12 +11027,7 @@ NV_STATUS uvm_va_block_find_create_in_range(uvm_va_space_t *va_space,
     if (uvm_enable_builtin_tests && atomic_dec_if_positive(&va_space->test.va_block_allocation_fail_nth) == 0)
         return NV_ERR_NO_MEMORY;
 
-    if (!va_range) {
-        if (!va_block_context || !va_block_context->mm)
-            return NV_ERR_INVALID_ADDRESS;
-        return uvm_hmm_va_block_find_create(va_space, addr, va_block_context, out_block);
-    }
-
+    UVM_ASSERT(va_range);
     UVM_ASSERT(addr >= va_range->node.start);
     UVM_ASSERT(addr <= va_range->node.end);
 
@@ -11057,14 +11041,32 @@ NV_STATUS uvm_va_block_find_create_in_range(uvm_va_space_t *va_space,
     return uvm_va_range_block_create(va_range, index, out_block);
 }
 
-NV_STATUS uvm_va_block_find_create(uvm_va_space_t *va_space,
+NV_STATUS uvm_va_block_find_create_managed(uvm_va_space_t *va_space,
                                    NvU64 addr,
-                                   uvm_va_block_context_t *va_block_context,
                                    uvm_va_block_t **out_block)
 {
     uvm_va_range_t *va_range = uvm_va_range_find(va_space, addr);
 
-    return uvm_va_block_find_create_in_range(va_space, va_range, addr, va_block_context, out_block);
+    if (va_range)
+        return uvm_va_block_find_create_in_range(va_space, va_range, addr, out_block);
+    else
+        return NV_ERR_INVALID_ADDRESS;
+}
+
+NV_STATUS uvm_va_block_find_create(uvm_va_space_t *va_space,
+                                   NvU64 addr,
+                                   struct vm_area_struct **hmm_vma,
+                                   uvm_va_block_t **out_block)
+{
+    uvm_va_range_t *va_range = uvm_va_range_find(va_space, addr);
+
+    if (hmm_vma)
+        *hmm_vma = NULL;
+
+    if (va_range)
+        return uvm_va_block_find_create_in_range(va_space, va_range, addr, out_block);
+    else
+        return uvm_hmm_va_block_find_create(va_space, addr, hmm_vma, out_block);
 }
 
 // Launch a synchronous, encrypted copy between GPU and CPU.
@@ -11241,8 +11243,6 @@ NV_STATUS uvm_va_block_write_from_cpu(uvm_va_block_t *va_block,
     if (UVM_ID_IS_INVALID(proc))
         proc = UVM_ID_CPU;
 
-    block_context->policy = uvm_va_policy_get(va_block, dst);
-
     // Use make_resident() in all cases to break read-duplication, but
     // block_retry can be NULL as if the page is not resident yet we will make
     // it resident on the CPU.
@@ -11411,7 +11411,6 @@ static void block_add_eviction_mappings(void *args)
         uvm_va_range_t *va_range = va_block->va_range;
         NV_STATUS status = NV_OK;
 
-        block_context->policy = uvm_va_range_get_policy(va_range);
         for_each_id_in_mask(id, &uvm_va_range_get_policy(va_range)->accessed_by) {
             status = uvm_va_block_set_accessed_by(va_block, block_context, id);
             if (status != NV_OK)
@@ -11520,6 +11519,11 @@ NV_STATUS uvm_va_block_evict_chunks(uvm_va_block_t *va_block,
         return NV_ERR_NO_MEMORY;
     }
 
+    if (uvm_va_block_is_hmm(va_block)) {
+        memset(block_context->hmm.src_pfns, 0, sizeof(block_context->hmm.src_pfns));
+        memset(block_context->hmm.dst_pfns, 0, sizeof(block_context->hmm.dst_pfns));
+    }
+
     pages_to_evict = &block_context->caller_page_mask;
     uvm_page_mask_zero(pages_to_evict);
     chunk_region.outer = 0;
@@ -11562,8 +11566,8 @@ NV_STATUS uvm_va_block_evict_chunks(uvm_va_block_t *va_block,
                                                &accessed_by_set);
     }
     else {
-        block_context->policy = uvm_va_range_get_policy(va_block->va_range);
-        accessed_by_set = uvm_processor_mask_get_count(&block_context->policy->accessed_by) > 0;
+        const uvm_va_policy_t *policy = uvm_va_range_get_policy(va_block->va_range);
+        accessed_by_set = uvm_processor_mask_get_count(&policy->accessed_by) > 0;
 
         // TODO: Bug 1765193: make_resident() breaks read-duplication, but it's
         // not necessary to do so for eviction. Add a version that unmaps only
@@ -11754,19 +11758,16 @@ NV_STATUS uvm_test_va_block_inject_error(UVM_TEST_VA_BLOCK_INJECT_ERROR_PARAMS *
     struct mm_struct *mm;
     uvm_va_block_t *va_block;
     uvm_va_block_test_t *va_block_test;
-    uvm_va_block_context_t *block_context = NULL;
     NV_STATUS status = NV_OK;
 
     mm = uvm_va_space_mm_or_current_retain_lock(va_space);
     uvm_va_space_down_read(va_space);
 
-    block_context = uvm_va_block_context_alloc(mm);
-    if (!block_context) {
-        status = NV_ERR_NO_MEMORY;
-        goto out;
-    }
+    if (mm)
+        status = uvm_va_block_find_create(va_space, params->lookup_address, NULL, &va_block);
+    else
+        status = uvm_va_block_find_create_managed(va_space, params->lookup_address, &va_block);
 
-    status = uvm_va_block_find_create(va_space, params->lookup_address, block_context, &va_block);
     if (status != NV_OK)
         goto out;
 
@@ -11806,7 +11807,6 @@ block_unlock:
 out:
     uvm_va_space_up_read(va_space);
     uvm_va_space_mm_or_current_release_unlock(va_space, mm);
-    uvm_va_block_context_free(block_context);
     return status;
 }
 
@@ -11877,7 +11877,11 @@ NV_STATUS uvm_test_change_pte_mapping(UVM_TEST_CHANGE_PTE_MAPPING_PARAMS *params
         goto out;
     }
 
-    status = uvm_va_block_find_create(va_space, params->va, block_context, &block);
+    if (mm)
+        status = uvm_va_block_find_create(va_space, params->va, &block_context->hmm.vma, &block);
+    else
+        status = uvm_va_block_find_create_managed(va_space, params->va, &block);
+
     if (status != NV_OK)
         goto out;
 
@@ -11904,8 +11908,6 @@ NV_STATUS uvm_test_change_pte_mapping(UVM_TEST_CHANGE_PTE_MAPPING_PARAMS *params
         goto out_block;
     }
 
-    block_context->policy = uvm_va_policy_get(block, params->va);
-
     if (new_prot == UVM_PROT_NONE) {
         status = uvm_va_block_unmap(block, block_context, id, region, NULL, &block->tracker);
     }

kernel-open/nvidia-uvm/uvm_va_block.h

@@ -453,11 +453,12 @@ struct uvm_va_block_struct
         NvU16 fault_migrations_to_last_proc;
     } prefetch_info;
 
-#if UVM_IS_CONFIG_HMM()
     struct
     {
+#if UVM_IS_CONFIG_HMM()
         // The MMU notifier is registered per va_block.
         struct mmu_interval_notifier notifier;
+#endif
 
         // This is used to serialize migrations between CPU and GPU while
         // allowing the va_block lock to be dropped.
@@ -487,7 +488,6 @@ struct uvm_va_block_struct
         // Storage node for range tree of va_blocks.
         uvm_range_tree_node_t node;
     } hmm;
-#endif
 };
 
 // We define additional per-VA Block fields for testing. When
@@ -678,18 +678,8 @@ static void uvm_va_block_context_init(uvm_va_block_context_t *va_block_context,
         memset(va_block_context, 0xff, sizeof(*va_block_context));
 
     va_block_context->mm = mm;
-#if UVM_IS_CONFIG_HMM()
-    va_block_context->hmm.vma = NULL;
-#endif
 }
 
-// Check that a single policy covers the given region for the given va_block.
-// This always returns true and is intended to only be used with UVM_ASSERT().
-// Locking: the va_block lock must be held.
-bool uvm_va_block_check_policy_is_valid(uvm_va_block_t *va_block,
-                                        const uvm_va_policy_t *policy,
-                                        uvm_va_block_region_t region);
-
 // TODO: Bug 1766480: Using only page masks instead of a combination of regions
 //       and page masks could simplify the below APIs and their implementations
 //       at the cost of having to scan the whole mask for small regions.
@@ -734,15 +724,15 @@ bool uvm_va_block_check_policy_is_valid(uvm_va_block_t *va_block,
 // user memory is guaranteed not to happen. Allocation-retry of GPU page tables
 // can still occur though.
 //
-// va_block_context must not be NULL. This function will set a bit in
-// va_block_context->make_resident.pages_changed_residency for each page that
-// changed residency (due to a migration or first population) as a result of the
-// operation and va_block_context->make_resident.all_involved_processors for
-// each processor involved in the copy. This function only sets bits in those
-// masks. It is the caller's responsiblity to zero the masks or not first.
-//
-// va_block_context->policy must also be set by the caller for the given region.
-// See the comments for uvm_va_block_check_policy_is_valid().
+// va_block_context must not be NULL and policy for the region must
+// match. This function will set a bit in
+// va_block_context->make_resident.pages_changed_residency for each
+// page that changed residency (due to a migration or first
+// population) as a result of the operation and
+// va_block_context->make_resident.all_involved_processors for each
+// processor involved in the copy. This function only sets bits in
+// those masks. It is the caller's responsiblity to zero the masks or
+// not first.
 //
 // Notably any status other than NV_OK indicates that the block's lock might
 // have been unlocked and relocked.
@@ -839,7 +829,7 @@ void uvm_va_block_make_resident_finish(uvm_va_block_t *va_block,
 // pages because the earlier operation can cause a PTE split or merge which is
 // assumed by the later operation.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // If allocation-retry was required as part of the operation and was successful,
@@ -896,7 +886,7 @@ NV_STATUS uvm_va_block_map_mask(uvm_va_block_t *va_block,
 // pages because the earlier operation can cause a PTE split or merge which is
 // assumed by the later operation.
 //
-// va_block_context must not be NULL. The va_block_context->policy is unused.
+// va_block_context must not be NULL.
 //
 // If allocation-retry was required as part of the operation and was successful,
 // NV_ERR_MORE_PROCESSING_REQUIRED is returned. In this case, the entries in the
@@ -929,7 +919,7 @@ NV_STATUS uvm_va_block_unmap_mask(uvm_va_block_t *va_block,
 // - Unmap the preferred location's processor from any pages in this region
 //   which are not resident on the preferred location.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: The caller must hold the VA block lock.
@@ -941,7 +931,7 @@ NV_STATUS uvm_va_block_set_preferred_location_locked(uvm_va_block_t *va_block,
 // location and policy. Waits for the operation to complete before returning.
 // This function should only be called with managed va_blocks.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: This takes and releases the VA block lock. If va_block_context->mm
@@ -956,7 +946,7 @@ NV_STATUS uvm_va_block_set_accessed_by(uvm_va_block_t *va_block,
 // the tracker after all mappings have been started.
 // This function can be called with HMM and managed va_blocks.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: The caller must hold the va_block lock and
@@ -970,7 +960,7 @@ NV_STATUS uvm_va_block_set_accessed_by_locked(uvm_va_block_t *va_block,
 // Breaks SetAccessedBy and remote mappings
 // This function should only be called with managed va_blocks.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: This takes and releases the VA block lock. If va_block_context->mm
@@ -982,7 +972,7 @@ NV_STATUS uvm_va_block_set_read_duplication(uvm_va_block_t *va_block,
 // Restores SetAccessedBy mappings
 // This function should only be called with managed va_blocks.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: This takes and releases the VA block lock. If va_block_context->mm
@@ -1002,16 +992,15 @@ NV_STATUS uvm_va_block_unset_read_duplication(uvm_va_block_t *va_block,
 // NV_ERR_INVALID_OPERATION     The access would violate the policies specified
 //                              by UvmPreventMigrationRangeGroups.
 //
-// va_block_context must not be NULL, va_block_context->policy must be valid,
-// and if the va_block is a HMM block, va_block_context->hmm.vma must be valid
-// which also means the va_block_context->mm is not NULL, retained, and locked
-// for at least read.
+// va_block_context must not be NULL, policy must match, and if the va_block is
+// a HMM block, va_block_context->hmm.vma must be valid which also means the
+// va_block_context->mm is not NULL, retained, and locked for at least read.
 // Locking: the va_block lock must be held.
 NV_STATUS uvm_va_block_check_logical_permissions(uvm_va_block_t *va_block,
                                                  uvm_va_block_context_t *va_block_context,
                                                  uvm_processor_id_t processor_id,
                                                  uvm_page_index_t page_index,
-                                                 uvm_fault_type_t access_type,
+                                                 uvm_fault_access_type_t access_type,
                                                  bool allow_migration);
 
 // API for access privilege revocation
@@ -1041,7 +1030,7 @@ NV_STATUS uvm_va_block_check_logical_permissions(uvm_va_block_t *va_block,
 // different pages because the earlier operation can cause a PTE split or merge
 // which is assumed by the later operation.
 //
-// va_block_context must not be NULL. The va_block_context->policy is unused.
+// va_block_context must not be NULL.
 //
 // If allocation-retry was required as part of the operation and was successful,
 // NV_ERR_MORE_PROCESSING_REQUIRED is returned. In this case, the entries in the
@@ -1081,7 +1070,7 @@ NV_STATUS uvm_va_block_revoke_prot_mask(uvm_va_block_t *va_block,
 // processor_id, which triggered the migration and should have already been
 // mapped).
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // This function acquires/waits for the va_block tracker and updates that
@@ -1112,7 +1101,7 @@ NV_STATUS uvm_va_block_add_mappings_after_migration(uvm_va_block_t *va_block,
 // Note that this can return NV_ERR_MORE_PROCESSING_REQUIRED just like
 // uvm_va_block_map() indicating that the operation needs to be retried.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
+// va_block_context must not be NULL and policy must for the region must match.
 // See the comments for uvm_va_block_check_policy_is_valid().
 //
 // LOCKING: The caller must hold the va block lock. If va_block_context->mm !=
@@ -1134,7 +1123,7 @@ NV_STATUS uvm_va_block_add_gpu_va_space(uvm_va_block_t *va_block, uvm_gpu_va_spa
 // If mm != NULL, that mm is used for any CPU mappings which may be created as
 // a result of this call. See uvm_va_block_context_t::mm for details.
 //
-// va_block_context must not be NULL. The va_block_context->policy is unused.
+// va_block_context must not be NULL.
 //
 // LOCKING: The caller must hold the va_block lock. If block_context->mm is not
 // NULL, the caller must hold mm->mmap_lock in at least read mode.
@@ -1225,7 +1214,6 @@ NV_STATUS uvm_va_block_split_locked(uvm_va_block_t *existing_va_block,
 //  - va_space lock must be held in at least read mode
 //
 // service_context->block_context.mm is ignored and vma->vm_mm is used instead.
-// service_context->block_context.policy is set by this function.
 //
 // Returns NV_ERR_INVALID_ACCESS_TYPE if a CPU mapping to fault_addr cannot be
 // accessed, for example because it's within a range group which is non-
@@ -1239,10 +1227,10 @@ NV_STATUS uvm_va_block_cpu_fault(uvm_va_block_t *va_block,
 // (migrations, cache invalidates, etc.) in response to the given service block
 // context.
 //
-// service_context must not be NULL and service_context->block_context.policy
-// must be valid. See the comments for uvm_va_block_check_policy_is_valid().
-// If va_block is a HMM block, va_block_context->hmm.vma must be valid.
-// See the comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
+// service_context must not be NULL and policy for service_context->region must
+// match. See the comments for uvm_va_block_check_policy_is_valid().  If
+// va_block is a HMM block, va_block_context->hmm.vma must be valid.  See the
+// comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
 // service_context->prefetch_hint is set by this function.
 //
 // Locking:
@@ -1267,10 +1255,10 @@ NV_STATUS uvm_va_block_service_locked(uvm_processor_id_t processor_id,
 // Performs population of the destination pages, unmapping and copying source
 // pages to new_residency.
 //
-// service_context must not be NULL and service_context->block_context.policy
-// must be valid. See the comments for uvm_va_block_check_policy_is_valid().
-// If va_block is a HMM block, va_block_context->hmm.vma must be valid.
-// See the comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
+// service_context must not be NULL and policy for service_context->region must
+// match.  See the comments for uvm_va_block_check_policy_is_valid().  If
+// va_block is a HMM block, va_block_context->hmm.vma must be valid.  See the
+// comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
 // service_context->prefetch_hint should be set before calling this function.
 //
 // Locking:
@@ -1296,10 +1284,10 @@ NV_STATUS uvm_va_block_service_copy(uvm_processor_id_t processor_id,
 // This updates the va_block residency state and maps the faulting processor_id
 // to the new residency (which may be remote).
 //
-// service_context must not be NULL and service_context->block_context.policy
-// must be valid. See the comments for uvm_va_block_check_policy_is_valid().
-// If va_block is a HMM block, va_block_context->hmm.vma must be valid.
-// See the comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
+// service_context must not be NULL and policy for service_context->region must
+// match. See the comments for uvm_va_block_check_policy_is_valid().  If
+// va_block is a HMM block, va_block_context->hmm.vma must be valid.  See the
+// comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
 // service_context must be initialized by calling uvm_va_block_service_copy()
 // before calling this function.
 //
@@ -1428,40 +1416,34 @@ const uvm_page_mask_t *uvm_va_block_map_mask_get(uvm_va_block_t *block, uvm_proc
 NV_STATUS uvm_va_block_find(uvm_va_space_t *va_space, NvU64 addr, uvm_va_block_t **out_block);
 
 // Same as uvm_va_block_find except that the block is created if not found.
-// If addr is covered by a UVM_VA_RANGE_TYPE_MANAGED va_range, a managed block
-// will be created. Otherwise, if addr is not covered by any va_range, HMM is
-// enabled in the va_space, and va_block_context and va_block_context->mm are
-// non-NULL, then a HMM block will be created and va_block_context->hmm.vma is
-// set to the VMA covering 'addr'. The va_block_context->policy field is left
-// unchanged.
-// In either case, if va_block_context->mm is non-NULL, it must be retained and
-// locked in at least read mode. Return values:
+// If addr is covered by a UVM_VA_RANGE_TYPE_MANAGED va_range a managed block
+// will be created. If addr is not covered by any va_range and HMM is
+// enabled in the va_space then a HMM block will be created and hmm_vma is
+// set to the VMA covering 'addr'. The va_space_mm must be retained and locked.
+// Otherwise hmm_vma is set to NULL.
+// Return values:
 // NV_ERR_INVALID_ADDRESS   addr is not a UVM_VA_RANGE_TYPE_MANAGED va_range nor
 //                          a HMM enabled VMA.
 // NV_ERR_NO_MEMORY         memory could not be allocated.
 NV_STATUS uvm_va_block_find_create(uvm_va_space_t *va_space,
                                    NvU64 addr,
-                                   uvm_va_block_context_t *va_block_context,
+                                   struct vm_area_struct **hmm_vma,
                                    uvm_va_block_t **out_block);
 
-// Same as uvm_va_block_find_create except that va_range lookup was already done
-// by the caller. If the supplied va_range is NULL, this function behaves just
-// like when the va_range lookup in uvm_va_block_find_create is NULL.
+// Same as uvm_va_block_find_create except that only managed va_blocks are
+// created if not already present in the VA range. Does not require va_space_mm
+// to be locked or retained.
+NV_STATUS uvm_va_block_find_create_managed(uvm_va_space_t *va_space,
+                                           NvU64 addr,
+                                           uvm_va_block_t **out_block);
+
+// Same as uvm_va_block_find_create_managed except that va_range lookup was
+// already done by the caller. The supplied va_range must not be NULL.
 NV_STATUS uvm_va_block_find_create_in_range(uvm_va_space_t *va_space,
                                             uvm_va_range_t *va_range,
                                             NvU64 addr,
-                                            uvm_va_block_context_t *va_block_context,
                                             uvm_va_block_t **out_block);
 
-// Same as uvm_va_block_find_create except that only managed va_blocks are
-// created if not already present in the VA range.
-static NV_STATUS uvm_va_block_find_create_managed(uvm_va_space_t *va_space,
-                                                  NvU64 addr,
-                                                  uvm_va_block_t **out_block)
-{
-    return uvm_va_block_find_create(va_space, addr, NULL, out_block);
-}
-
 // Look up a chunk backing a specific address within the VA block.
 // Returns NULL if none.
 uvm_gpu_chunk_t *uvm_va_block_lookup_gpu_chunk(uvm_va_block_t *va_block, uvm_gpu_t *gpu, NvU64 address);
@@ -1476,10 +1458,10 @@ uvm_gpu_chunk_t *uvm_va_block_lookup_gpu_chunk(uvm_va_block_t *va_block, uvm_gpu
 // The caller needs to handle allocation-retry. va_block_retry can be NULL if
 // the destination is the CPU.
 //
-// va_block_context must not be NULL and va_block_context->policy must be valid.
-// See the comments for uvm_va_block_check_policy_is_valid().
-// If va_block is a HMM block, va_block_context->hmm.vma must be valid.
-// See the comments for uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
+// va_block_context must not be NULL and policy for the region must match. See
+// the comments for uvm_va_block_check_policy_is_valid().  If va_block is a HMM
+// block, va_block_context->hmm.vma must be valid.  See the comments for
+// uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
 //
 // LOCKING: The caller must hold the va_block lock. If va_block_context->mm !=
 //          NULL, va_block_context->mm->mmap_lock must be held in at least
@@ -1497,7 +1479,7 @@ NV_STATUS uvm_va_block_migrate_locked(uvm_va_block_t *va_block,
 // The [dst, dst + size) range has to fit within a single PAGE_SIZE page.
 //
 // va_block_context must not be NULL. The caller is not required to set
-// va_block_context->policy or va_block_context->hmm.vma.
+// va_block_context->hmm.vma.
 //
 // The caller needs to support allocation-retry of page tables.
 //
@@ -1569,7 +1551,7 @@ void uvm_va_block_mark_cpu_dirty(uvm_va_block_t *va_block);
 // successful, NV_ERR_MORE_PROCESSING_REQUIRED is returned. In this case the
 // block's lock was unlocked and relocked.
 //
-// va_block_context must not be NULL. The va_block_context->policy is unused.
+// va_block_context must not be NULL.
 //
 // LOCKING: The caller must hold the va_block lock.
 NV_STATUS uvm_va_block_set_cancel(uvm_va_block_t *va_block, uvm_va_block_context_t *block_context, uvm_gpu_t *gpu);
@@ -1650,12 +1632,18 @@ static uvm_va_block_region_t uvm_va_block_region_from_block(uvm_va_block_t *va_b
     return uvm_va_block_region(0, uvm_va_block_num_cpu_pages(va_block));
 }
 
-// Create a block region from a va block and page mask. Note that the region
+// Create a block region from a va block and page mask. If va_block is NULL, the
+// region is assumed to cover the maximum va_block size. Note that the region
 // covers the first through the last set bit and may have unset bits in between.
 static uvm_va_block_region_t uvm_va_block_region_from_mask(uvm_va_block_t *va_block, const uvm_page_mask_t *page_mask)
 {
     uvm_va_block_region_t region;
-    uvm_page_index_t outer = uvm_va_block_num_cpu_pages(va_block);
+    uvm_page_index_t outer;
+
+    if (va_block)
+        outer = uvm_va_block_num_cpu_pages(va_block);
+    else
+        outer = PAGES_PER_UVM_VA_BLOCK;
 
     region.first = find_first_bit(page_mask->bitmap, outer);
     if (region.first >= outer) {
@@ -2084,6 +2072,14 @@ void uvm_va_block_unmap_cpu_chunk_on_gpus(uvm_va_block_t *va_block,
 // Locking: The va_block lock must be held.
 void uvm_va_block_remove_cpu_chunks(uvm_va_block_t *va_block, uvm_va_block_region_t region);
 
+// Get the size of the physical allocation backing the page at page_index on the
+// specified processor in the block. Returns 0 if the address is not resident on
+// the specified processor.
+// Locking: The va_block lock must be held.
+NvU32 uvm_va_block_get_physical_size(uvm_va_block_t *block,
+                                     uvm_processor_id_t processor,
+                                     uvm_page_index_t page_index);
+
 // Get CPU page size or 0 if it is not mapped
 NvU32 uvm_va_block_page_size_cpu(uvm_va_block_t *va_block,
                                  uvm_page_index_t page_index);
@@ -2140,15 +2136,14 @@ uvm_va_block_region_t uvm_va_block_big_page_region_subset(uvm_va_block_t *va_blo
 // MAX_BIG_PAGES_PER_UVM_VA_BLOCK is returned.
 size_t uvm_va_block_big_page_index(uvm_va_block_t *va_block, uvm_page_index_t page_index, NvU32 big_page_size);
 
-// Returns the new residency for a page that faulted or triggered access
-// counter notifications. The read_duplicate output parameter indicates if the
-// page meets the requirements to be read-duplicated
-// va_block_context must not be NULL, va_block_context->policy must be valid,
-// and if the va_block is a HMM block, va_block_context->hmm.vma must be valid
-// which also means the va_block_context->mm is not NULL, retained, and locked
-// for at least read. See the comments for uvm_va_block_check_policy_is_valid()
-// and uvm_hmm_check_context_vma_is_valid() in uvm_hmm.h.
-// Locking: the va_block lock must be held.
+// Returns the new residency for a page that faulted or triggered access counter
+// notifications. The read_duplicate output parameter indicates if the page
+// meets the requirements to be read-duplicated va_block_context must not be
+// NULL, and if the va_block is a HMM block, va_block_context->hmm.vma must be
+// valid which also means the va_block_context->mm is not NULL, retained, and
+// locked for at least read. See the comments for
+// uvm_va_block_check_policy_is_valid() and uvm_hmm_check_context_vma_is_valid()
+// in uvm_hmm.h.  Locking: the va_block lock must be held.
 uvm_processor_id_t uvm_va_block_select_residency(uvm_va_block_t *va_block,
                                                  uvm_va_block_context_t *va_block_context,
                                                  uvm_page_index_t page_index,

kernel-open/nvidia-uvm/uvm_va_block_types.h

@@ -29,9 +29,7 @@
 #include "uvm_tlb_batch.h"
 #include "uvm_forward_decl.h"
 
-#if UVM_IS_CONFIG_HMM()
 #include <linux/migrate.h>
-#endif
 
 // UVM_VA_BLOCK_BITS is 21, meaning the maximum block size is 2MB. Rationale:
 // - 2MB matches the largest Pascal GPU page size so it's a natural fit
@@ -234,9 +232,6 @@ typedef struct
     // the mm, such as creating CPU mappings.
     struct mm_struct *mm;
 
-    const uvm_va_policy_t *policy;
-
-#if UVM_IS_CONFIG_HMM()
     struct
     {
         // These are used for migrate_vma_*(), hmm_range_fault(), and
@@ -257,10 +252,11 @@ typedef struct
         // Cached VMA pointer. This is only valid while holding the mmap_lock.
         struct vm_area_struct *vma;
 
+#if UVM_IS_CONFIG_HMM()
         // Used for migrate_vma_*() to migrate pages to/from GPU/CPU.
         struct migrate_vma migrate_vma_args;
-    } hmm;
 #endif
+    } hmm;
 
     // Convenience buffer for page mask prints
     char page_mask_string_buffer[UVM_PAGE_MASK_PRINT_MIN_BUFFER_SIZE];

kernel-open/nvidia-uvm/uvm_va_policy.c

@@ -54,6 +54,52 @@ const uvm_va_policy_t *uvm_va_policy_get(uvm_va_block_t *va_block, NvU64 addr)
     }
 }
 
+// HMM va_blocks can have different polices for different regions withing the
+// va_block. This function checks the given region is covered by the same policy
+// and asserts if the region is covered by different policies.
+// This always returns true and is intended to only be used with UVM_ASSERT() to
+// avoid calling it on release builds.
+// Locking: the va_block lock must be held.
+static bool uvm_hmm_va_block_assert_policy_is_valid(uvm_va_block_t *va_block,
+                                                    const uvm_va_policy_t *policy,
+                                                    uvm_va_block_region_t region)
+{
+    const uvm_va_policy_node_t *node;
+
+    if (uvm_va_policy_is_default(policy)) {
+        // There should only be the default policy within the region.
+        node = uvm_va_policy_node_iter_first(va_block,
+                                             uvm_va_block_region_start(va_block, region),
+                                             uvm_va_block_region_end(va_block, region));
+        UVM_ASSERT(!node);
+    }
+    else {
+        // The policy node should cover the region.
+        node = uvm_va_policy_node_from_policy(policy);
+        UVM_ASSERT(node->node.start <= uvm_va_block_region_start(va_block, region));
+        UVM_ASSERT(node->node.end >= uvm_va_block_region_end(va_block, region));
+    }
+
+    return true;
+}
+
+const uvm_va_policy_t *uvm_va_policy_get_region(uvm_va_block_t *va_block, uvm_va_block_region_t region)
+{
+    uvm_assert_mutex_locked(&va_block->lock);
+
+    if (uvm_va_block_is_hmm(va_block)) {
+        const uvm_va_policy_t *policy;
+        const uvm_va_policy_node_t *node = uvm_va_policy_node_find(va_block, uvm_va_block_region_start(va_block, region));
+
+        policy = node ? &node->policy : &uvm_va_policy_default;
+        UVM_ASSERT(uvm_hmm_va_block_assert_policy_is_valid(va_block, policy, region));
+        return policy;
+    }
+    else {
+        return uvm_va_range_get_policy(va_block->va_range);
+    }
+}
+
 #if UVM_IS_CONFIG_HMM()
 
 static struct kmem_cache *g_uvm_va_policy_node_cache __read_mostly;