Fix cudaMallocManaged() on nvproxy.

- The app attempts to allocate a driver object of class NV_CONFIDENTIAL_COMPUTE
  in drivers too old to support said class; said allocation must fail with
  status NV_ERR_INVALID_CLASS rather than errno EINVAL for the app to proceed.

- UVM_VALIDATE_VA_RANGE checks that a given address range is known to
  nvidia-uvm. For mmaps of /dev/nvidia-uvm, this requires that the application
  mmap (handled by nvproxy) immediately result in a host mmap (handled by the
  driver). Ensure that mappings of nvproxy's nvidia-uvm have this property.

- Pass through UVM ioctl UVM_DISABLE_READ_DUPLICATION.

Fixes #9593

PiperOrigin-RevId: 577966817

Makefile

@@ -262,7 +262,7 @@ arm-qemu-smoke-test: $(RUNTIME_BIN) load-arm-qemu
 simple-tests: unit-tests # Compatibility target.
 .PHONY: simple-tests
 
-gpu-tests: load-basic $(RUNTIME_BIN)
+gpu-tests: load-basic_cuda-vector-add load-gpu_cuda-tests $(RUNTIME_BIN)
 	@$(call test,--test_env=RUNTIME=runc //test/gpu:gpu_test)
 	@$(call install_runtime,$(RUNTIME),--platform=systrap --nvproxy=true --nvproxy-docker=true)
 	@$(call test_runtime,$(RUNTIME),//test/gpu:gpu_test)

images/gpu/cuda-tests/Dockerfile

@@ -0,0 +1,7 @@
+FROM nvidia/cuda:12.2.0-devel-ubuntu20.04
+
+WORKDIR /
+COPY cuda_malloc_managed.cu .
+COPY cuda_test_util.h .
+COPY run.sh .
+ENTRYPOINT ["/run.sh"]

images/gpu/cuda-tests/cuda_malloc_managed.cu

@@ -0,0 +1,101 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <iostream>
+#include <random>
+
+#include "cuda_test_util.h"  // NOLINT(build/include)
+
+__global__ void addKernel(std::uint32_t* data) {
+  size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  data[index] += static_cast<std::uint32_t>(index);
+}
+
+void TestMallocManagedRoundTrip(int device, unsigned int malloc_flags,
+                                bool prefetch) {
+  constexpr size_t kNumBlocks = 32;
+  constexpr size_t kNumThreads = 64;
+  constexpr size_t kNumElems = kNumBlocks * kNumThreads;
+
+  std::uint32_t* data = nullptr;
+  constexpr size_t kNumBytes = kNumElems * sizeof(*data);
+  CHECK_CUDA(cudaMallocManaged(&data, kNumBytes, malloc_flags));
+
+  // Initialize all elements in the array with a random value on the host.
+  std::random_device rd;
+  const std::uint32_t init_val =
+      std::uniform_int_distribution<std::uint32_t>()(rd);
+  for (size_t i = 0; i < kNumElems; i++) {
+    data[i] = init_val;
+  }
+
+  if (prefetch) {
+    CHECK_CUDA(cudaMemPrefetchAsync(data, kNumBytes, device));
+  }
+
+  // Mutate the array on the device.
+  addKernel<<<kNumBlocks, kNumThreads>>>(data);
+  CHECK_CUDA(cudaDeviceSynchronize());
+
+  if (prefetch) {
+    CHECK_CUDA(cudaMemPrefetchAsync(data, kNumBytes, cudaCpuDeviceId));
+  }
+
+  // Check that the array has the expected result.
+  for (size_t i = 0; i < kNumElems; i++) {
+    std::uint32_t want = init_val + static_cast<std::uint32_t>(i);
+    if (data[i] != want) {
+      std::cout << "data[" << i << "]: got " << data[i] << ", wanted " << want
+                << " = " << init_val << " + " << i << std::endl;
+      abort();
+    }
+  }
+
+  CHECK_CUDA(cudaFree(data));
+}
+
+int main() {
+  int device;
+  CHECK_CUDA(cudaGetDevice(&device));
+
+  std::cout << "Testing cudaMallocManaged(flags=cudaMemAttachGlobal)"
+            << std::endl;
+  TestMallocManagedRoundTrip(device, cudaMemAttachGlobal, false);
+
+  int cma = 0;
+  CHECK_CUDA(
+      cudaDeviceGetAttribute(&cma, cudaDevAttrConcurrentManagedAccess, device));
+  if (!cma) {
+    std::cout << "cudaDevAttrConcurrentManagedAccess not available"
+              << std::endl;
+  } else {
+    std::cout << "Testing cudaMallocManaged(flags=cudaMemAttachGlobal) "
+                 "with prefetching"
+              << std::endl;
+    TestMallocManagedRoundTrip(device, cudaMemAttachGlobal, true);
+    std::cout << "Testing cudaMallocManaged(flags=cudaMemAttachHost)"
+              << std::endl;
+    TestMallocManagedRoundTrip(device, cudaMemAttachHost, false);
+    std::cout << "Testing cudaMallocManaged(flags=cudaMemAttachHost) "
+                 "with prefetching"
+              << std::endl;
+    TestMallocManagedRoundTrip(device, cudaMemAttachHost, true);
+  }
+
+  std::cout << "All tests passed" << std::endl;
+  return 0;
+}

images/gpu/cuda-tests/cuda_test_util.h

@@ -0,0 +1,30 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef THIRD_PARTY_GVISOR_IMAGES_GPU_CUDA_TESTS_CUDA_TEST_UTIL_H_
+#define THIRD_PARTY_GVISOR_IMAGES_GPU_CUDA_TESTS_CUDA_TEST_UTIL_H_
+
+#include <iostream>
+
+#define CHECK_CUDA(expr)                                                     \
+  do {                                                                       \
+    cudaError_t code = (expr);                                               \
+    if (code != cudaSuccess) {                                               \
+      std::cout << "Check failed at " << __FILE__ << ":" << __LINE__ << ": " \
+                << #expr << ": " << cudaGetErrorString(code) << std::endl;   \
+      abort();                                                               \
+    }                                                                        \
+  } while (0)
+
+#endif  // THIRD_PARTY_GVISOR_IMAGES_GPU_CUDA_TESTS_CUDA_TEST_UTIL_H_

images/gpu/cuda-tests/run.sh

@@ -0,0 +1,21 @@
+#!/bin/sh
+
+# Copyright 2021 The gVisor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eux
+
+cd /
+nvcc cuda_malloc_managed.cu -o cuda_malloc_managed
+./cuda_malloc_managed

pkg/abi/nvgpu/status.go

@@ -18,6 +18,7 @@ package nvgpu
 const (
 	NV_ERR_INVALID_ADDRESS  = 0x0000001e
 	NV_ERR_INVALID_ARGUMENT = 0x0000001f
+	NV_ERR_INVALID_CLASS    = 0x00000022
 	NV_ERR_INVALID_LIMIT    = 0x0000002e
 	NV_ERR_NOT_SUPPORTED    = 0x00000056
 )

pkg/abi/nvgpu/uvm.go

@@ -40,6 +40,7 @@ const (
 	UVM_REGISTER_GPU                   = 37
 	UVM_UNREGISTER_GPU                 = 38
 	UVM_PAGEABLE_MEM_ACCESS            = 39
+	UVM_DISABLE_READ_DUPLICATION       = 45
 	UVM_MAP_DYNAMIC_PARALLELISM_REGION = 65
 	UVM_ALLOC_SEMAPHORE_POOL           = 68
 	UVM_VALIDATE_VA_RANGE              = 72
@@ -187,6 +188,14 @@ type UVM_PAGEABLE_MEM_ACCESS_PARAMS struct {
 	RMStatus          uint32
 }
 
+// +marshal
+type UVM_DISABLE_READ_DUPLICATION_PARAMS struct {
+	RequestedBase uint64
+	Length        uint64
+	RMStatus      uint32
+	Pad0          [4]byte
+}
+
 // +marshal
 type UVM_MAP_DYNAMIC_PARALLELISM_REGION_PARAMS struct {
 	Base     uint64

pkg/sentry/devices/nvproxy/frontend.go

@@ -648,7 +648,17 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 	handler := fi.fd.nvp.abi.allocationClass[ioctlParams.HClass]
 	if handler == nil {
 		fi.ctx.Warningf("nvproxy: unknown allocation class %#08x", ioctlParams.HClass)
-		return 0, linuxerr.EINVAL
+		// Compare
+		// src/nvidia/src/kernel/rmapi/alloc_free.c:serverAllocResourceUnderLock(),
+		// when RsResInfoByExternalClassId() is null.
+		ioctlParams.Status = nvgpu.NV_ERR_INVALID_CLASS
+		outIoctlParams := nvgpu.GetRmAllocParamObj(isNVOS64, fi.fd.nvp.abi.useRmAllocParamsV535)
+		outIoctlParams.FromOS64V535(ioctlParams)
+		// Any copy-out error from
+		// src/nvidia/src/kernel/rmapi/alloc_free.c:serverAllocApiCopyOut() is
+		// discarded.
+		outIoctlParams.CopyOut(fi.t, fi.ioctlParamsAddr)
+		return 0, nil
 	}
 	return handler(fi, &ioctlParams, isNVOS64)
 }

pkg/sentry/devices/nvproxy/seccomp_filters.go

@@ -167,6 +167,10 @@ func Filters() seccomp.SyscallRules {
 				nonNegativeFD,
 				seccomp.EqualTo(nvgpu.UVM_PAGEABLE_MEM_ACCESS),
 			},
+			seccomp.PerArg{
+				nonNegativeFD,
+				seccomp.EqualTo(nvgpu.UVM_DISABLE_READ_DUPLICATION),
+			},
 			seccomp.PerArg{
 				nonNegativeFD,
 				seccomp.EqualTo(nvgpu.UVM_MAP_DYNAMIC_PARALLELISM_REGION),

pkg/sentry/devices/nvproxy/uvm.go

@@ -126,7 +126,7 @@ func (fd *uvmFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args
 	}
 
 	if log.IsLogging(log.Debug) {
-		ctx.Debugf("nvproxy: frontend ioctl %#08x", cmd)
+		ctx.Debugf("nvproxy: uvm ioctl %#08x", cmd)
 	}
 
 	ui := uvmIoctlState{

pkg/sentry/devices/nvproxy/uvm_mmap.go

@@ -26,6 +26,12 @@ import (
 
 // ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
 func (fd *uvmFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
+	// UVM_VALIDATE_VA_RANGE, and probably other ioctls, expect that
+	// application mmaps of /dev/nvidia-uvm are immediately visible to the
+	// driver.
+	if opts.PlatformEffect < memmap.PlatformEffectPopulate {
+		opts.PlatformEffect = memmap.PlatformEffectPopulate
+	}
 	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
 }
 

pkg/sentry/devices/nvproxy/version.go

@@ -188,6 +188,7 @@ func Init() {
 					nvgpu.UVM_REGISTER_GPU:                   uvmIoctlHasRMCtrlFD[nvgpu.UVM_REGISTER_GPU_PARAMS],
 					nvgpu.UVM_UNREGISTER_GPU:                 uvmIoctlSimple[nvgpu.UVM_UNREGISTER_GPU_PARAMS],
 					nvgpu.UVM_PAGEABLE_MEM_ACCESS:            uvmIoctlSimple[nvgpu.UVM_PAGEABLE_MEM_ACCESS_PARAMS],
+					nvgpu.UVM_DISABLE_READ_DUPLICATION:       uvmIoctlSimple[nvgpu.UVM_DISABLE_READ_DUPLICATION_PARAMS],
 					nvgpu.UVM_MAP_DYNAMIC_PARALLELISM_REGION: uvmIoctlSimple[nvgpu.UVM_MAP_DYNAMIC_PARALLELISM_REGION_PARAMS],
 					nvgpu.UVM_ALLOC_SEMAPHORE_POOL:           uvmIoctlSimple[nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS],
 					nvgpu.UVM_VALIDATE_VA_RANGE:              uvmIoctlSimple[nvgpu.UVM_VALIDATE_VA_RANGE_PARAMS],

test/gpu/gpu_test.go

@@ -41,8 +41,31 @@ func TestGPUHello(t *testing.T) {
 	})
 
 	if err != nil {
-		t.Fatalf("could not run nvidia: %v", err)
+		t.Fatalf("could not run cuda-vector-add: %v", err)
 	}
 
-	t.Logf("nvidia output: %s", string(out))
+	t.Logf("cuda-vector-add output: %s", string(out))
+}
+
+func TestCUDATests(t *testing.T) {
+	ctx := context.Background()
+	c := dockerutil.MakeContainer(ctx, t)
+	defer c.CleanUp(ctx)
+
+	out, err := c.Run(ctx, dockerutil.RunOpts{
+		Image: "gpu/cuda-tests",
+		Devices: []container.DeviceRequest{
+			{
+				Count:        -1,
+				Capabilities: [][]string{[]string{"gpu"}},
+				Options:      map[string]string{},
+			},
+		},
+	})
+
+	if err != nil {
+		t.Fatalf("could not run cuda-tests: %v", err)
+	}
+
+	t.Logf("cuda-tests output: %s", string(out))
 }