1 /*******************************************************************************
2     Copyright (c) 2018-2022 NVIDIA Corporation
3 
4     Permission is hereby granted, free of charge, to any person obtaining a copy
5     of this software and associated documentation files (the "Software"), to
6     deal in the Software without restriction, including without limitation the
7     rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
8     sell copies of the Software, and to permit persons to whom the Software is
9     furnished to do so, subject to the following conditions:
10 
11         The above copyright notice and this permission notice shall be
12         included in all copies or substantial portions of the Software.
13 
14     THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15     IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16     FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
17     THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18     LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
19     FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
20     DEALINGS IN THE SOFTWARE.
21 
22 *******************************************************************************/
23 
24 #include "uvm_common.h"
25 #include "uvm_kvmalloc.h"
26 #include "uvm_va_space.h"
27 #include "uvm_va_space_mm.h"
28 #include "uvm_ats.h"
29 #include "uvm_api.h"
30 #include "uvm_test.h"
31 #include "uvm_test_ioctl.h"
32 
33 //
34 // This comment block describes some implementation rationale. See the header
35 // for the API descriptions.
36 //
37 // ========================= Retain count vs mm_users ==========================
38 //
39 // To guarantee the mm is available and won't be destroyed we require
40 // userspace to open a second file descriptor (uvm_mm_fd) and
41 // initialize it with uvm_api_mm_initialize(). During initialization
42 // we take a mm_users reference to ensure the mm remains valid until
43 // the file descriptor is closed.
44 //
45 // To ensure userspace can't close the file descriptor and drop the
46 // mm_users refcount while it is in use threads must call either
47 // uvm_va_space_mm_retain() or uvm_va_space_mm_or_current_retain() to
48 // increment the retained count. This also checks that userspace has
49 // initialized the uvm_mm_fd and therefore holds a valid pagetable
50 // pin.
51 //
52 // Closing uvm_mm_fd will call uvm_va_space_mm_shutdown() prior to
53 // mmput() which ensures there are no active users of the mm. This
54 // indirection is required because not all threads can call mmput()
55 // directly. In particular the replayable GPU fault handling path
56 // can't call mmput() because it may result in exit_mmap() which could
57 // result in RM calls and VA space destroy and those need to wait for
58 // the GPU fault handler to finish.
59 //
60 // ============================ Handling mm teardown ===========================
61 //
62 // When the process is exiting we will get notified either via an
63 // explict close of uvm_mm_fd or implicitly as part of
64 // exit_files(). We are guaranteed to get this call because we don't
65 // allow mmap on uvm_mm_fd, and the userspace pagetables (mm_users)
66 // are guaranteed to exist because we hold a mm_users refcount
67 // which is released as part of file close.
68 //
69 // This allows any outstanding GPU faults to be processed. To prevent
70 // new faults occurring uvm_va_space_mm_shutdown() is called to stop
71 // all GPU memory accesses to the mm. Once all GPU memory has been
72 // stopped no new retainers of the va_space will be allowed and the
73 // mm_users reference will be dropped, potentially tearing down the mm
74 // and associated pagetables.
75 //
76 // This essentially shuts down the VA space for new work. The VA space
77 // object remains valid for most teardown ioctls until the file is
78 // closed, because it's legal for the associated process to die then
79 // for another process with a reference on the file to perform the
80 // unregisters or associated ioctls.  This is particularly true for
81 // tools users.
82 //
83 // An exception to the above is UvmUnregisterChannel. Since channels are
84 // completely removed from the VA space on mm teardown, later channel
85 // unregisters will fail to find the handles and will return an error.
86 //
87 // At a high level, the sequence of operations to perform prior to mm
88 // teardown is:
89 //
90 // 1) Stop all channels
91 //      - Prevents new faults and accesses on non-MPS
92 // 2) Detach all channels
93 //      - Prevents pending faults from being translated to this VA space
94 //      - Non-replayable faults will be dropped so no new ones can arrive
95 //      - Access counter notifications will be prevented from getting new
96 //        translations to this VA space. Pending entries may attempt to retain
97 //        the mm, but will drop the notification if they can't be serviced.
98 // 3) Flush the fault buffer
99 //      - The only reason to flush the fault buffer is to avoid spurious
100 //        cancels. If we didn't flush the fault buffer before marking the mm
101 //        as dead, then remaining faults which require the mm would be
102 //        cancelled. Since the faults might be stale, we would record cancel
103 //        events which didn't really happen (the access didn't happen after
104 //        the mm died). By flushing we clear out all stale faults, and in
105 //        the case of MPS, cancel real faults after.
106 // 4) UnsetPageDir
107 //      - Prevents new accesses on MPS
108 // 5) Mark the va_space_mm as released
109 //      - Prevents new retainers from using the mm. There won't be any more on
110 //        the fault handling paths, but there could be others in worker threads.
111 //
112 // Here are some tables of each step in the sequence, and what operations can
113 // still be performed after each step. This is all from the perspective of a
114 // single VA space. "Untranslated" means that the fault entry has not been
115 // translated to a uvm_va_space yet.
116 //
117 // Replayable non-MPS Behavior:
118 //
119 //                  Can              Pending         Pending         Can be
120 //                  access   Can     untranslated    translated      servicing
121 //                  memory   fault   faults          faults          faults
122 // -----------------------------------------------------------------------------
123 // Shutdown start   Yes      Yes     Service         Service         Yes
124 // Stop channels    No       No      Service [1]     Service [1]     Yes [1]
125 // Detach channels  No       No      Flush buffer    Service [1]     Yes [1], [2]
126 // Flush buffer     No       No      None possible   None possible   No
127 // UnsetPageDir     No       No      None possible   None possible   No
128 //
129 //
130 // Replayable MPS Behavior:
131 //
132 //                  Can              Pending         Pending         Can be
133 //                  access   Can     untranslated    translated      servicing
134 //                  memory   fault   faults          faults          faults
135 // -----------------------------------------------------------------------------
136 // Shutdown start   Yes      Yes     Service         Service         Yes
137 // Stop channels    Yes      Yes     Service         Service         Yes
138 // Detach channels  Yes      Yes     Cancel, flush   Service         Yes
139 // Flush buffer     Yes      Yes     Cancel, flush   None possible   No
140 // UnsetPageDir     No [3]   Yes     Cancel, flush   None possible   No
141 //
142 //
143 // [1]: All pending faults in this VA space are stale since channel stop
144 //      preempted the context.
145 // [2]: Faults in this VA space can't be serviced concurrently with detach since
146 //      detach holds the VA space lock in write mode. Faults in other VA spaces
147 //      can be serviced, and stale faults in this VA space can resume service
148 //      after detach is done.
149 // [3]: Due to the nature of MPS, remaining work which had started under the VA
150 //      space could still execute and attempt to make memory accesses. However,
151 //      since the PDB at that point is empty and ATS is disabled (if available),
152 //      all accesses will fault and be cancelled rather than successfully
153 //      translate to physical memory.
154 //
155 // =============================================================================
156 
157 static void uvm_va_space_mm_shutdown(uvm_va_space_t *va_space);
158 
159 static int uvm_enable_va_space_mm = 1;
160 module_param(uvm_enable_va_space_mm, int, S_IRUGO);
161 MODULE_PARM_DESC(uvm_enable_va_space_mm,
162                  "Set to 0 to disable UVM from using mmu_notifiers to create "
163                  "an association between a UVM VA space and a process. This "
164                  "will also disable pageable memory access via either ATS or "
165                  "HMM.");
166 
167 bool uvm_va_space_mm_enabled_system(void)
168 {
169     return UVM_CAN_USE_MMU_NOTIFIERS() && uvm_enable_va_space_mm;
170 }
171 
172 bool uvm_va_space_mm_enabled(uvm_va_space_t *va_space)
173 {
174     // A va_space doesn't have any association with an mm in multi-process
175     // sharing mode.
176     if (va_space->initialization_flags & UVM_INIT_FLAGS_MULTI_PROCESS_SHARING_MODE)
177         return false;
178 
179     return uvm_va_space_mm_enabled_system();
180 }
181 
182 #if UVM_CAN_USE_MMU_NOTIFIERS()
183     static uvm_va_space_t *get_va_space(struct mmu_notifier *mn)
184     {
185         // This may be called without a thread context present, so be careful
186         // what is used here.
187         return container_of(mn, uvm_va_space_t, va_space_mm.mmu_notifier);
188     }
189 
190     static void uvm_mmu_notifier_invalidate_range_ats(struct mmu_notifier *mn,
191                                                       struct mm_struct *mm,
192                                                       unsigned long start,
193                                                       unsigned long end)
194     {
195         // In most cases ->invalidate_range() is called with exclusive end.
196         // uvm_ats_invalidate() expects an inclusive end so we have to
197         // convert it.
198         //
199         // There's a special case however. Kernel TLB gathering sometimes
200         // identifies "fullmm" invalidates by setting both start and end to ~0.
201         //
202         // It's unclear if there are any other cases in which the kernel will
203         // call us with start == end. Since we can't definitively say no, we
204         // conservatively treat all such calls as full invalidates.
205         if (start == end) {
206             start = 0;
207             end = ~0UL;
208         }
209         else {
210             --end;
211         }
212 
213         UVM_ENTRY_VOID(uvm_ats_invalidate(get_va_space(mn), start, end));
214     }
215 
216     static struct mmu_notifier_ops uvm_mmu_notifier_ops_ats =
217     {
218         .invalidate_range = uvm_mmu_notifier_invalidate_range_ats,
219     };
220 
221     static int uvm_mmu_notifier_register(uvm_va_space_mm_t *va_space_mm)
222     {
223         UVM_ASSERT(va_space_mm->mm);
224         uvm_assert_mmap_lock_locked_write(va_space_mm->mm);
225 
226         va_space_mm->mmu_notifier.ops = &uvm_mmu_notifier_ops_ats;
227         return __mmu_notifier_register(&va_space_mm->mmu_notifier, va_space_mm->mm);
228     }
229 
230     static void uvm_mmu_notifier_unregister(uvm_va_space_mm_t *va_space_mm)
231     {
232         mmu_notifier_unregister(&va_space_mm->mmu_notifier, va_space_mm->mm);
233     }
234 #else
235     static int uvm_mmu_notifier_register(uvm_va_space_mm_t *va_space_mm)
236     {
237         UVM_ASSERT(0);
238         return 0;
239     }
240 
241     static void uvm_mmu_notifier_unregister(uvm_va_space_mm_t *va_space_mm)
242     {
243         UVM_ASSERT(0);
244     }
245 #endif // UVM_CAN_USE_MMU_NOTIFIERS()
246 
247 NV_STATUS uvm_va_space_mm_register(uvm_va_space_t *va_space)
248 {
249     uvm_va_space_mm_t *va_space_mm = &va_space->va_space_mm;
250     int ret;
251 
252     uvm_assert_mmap_lock_locked_write(current->mm);
253     uvm_assert_rwsem_locked_write(&va_space->lock);
254 
255     va_space_mm->state = UVM_VA_SPACE_MM_STATE_UNINITIALIZED;
256 
257     if (!uvm_va_space_mm_enabled(va_space))
258         return NV_OK;
259 
260     UVM_ASSERT(!va_space_mm->mm);
261     va_space_mm->mm = current->mm;
262     uvm_mmgrab(va_space_mm->mm);
263 
264     // We must be prepared to handle callbacks as soon as we make this call,
265     // except for ->release() which can't be called since the mm belongs to
266     // current.
267     if (UVM_ATS_IBM_SUPPORTED_IN_DRIVER() && g_uvm_global.ats.enabled) {
268         ret = uvm_mmu_notifier_register(va_space_mm);
269         if (ret) {
270             // Inform uvm_va_space_mm_unregister() that it has nothing to do.
271             uvm_mmdrop(va_space_mm->mm);
272             va_space_mm->mm = NULL;
273             return errno_to_nv_status(ret);
274         }
275     }
276 
277     if ((UVM_IS_CONFIG_HMM() || UVM_ATS_PREFETCH_SUPPORTED()) && uvm_va_space_pageable_mem_access_supported(va_space)) {
278         #if UVM_CAN_USE_MMU_NOTIFIERS()
279             // Initialize MMU interval notifiers for this process. This allows
280             // mmu_interval_notifier_insert() to be called without holding the
281             // mmap_lock for write.
282             // Note: there is no __mmu_notifier_unregister(), this call just
283             // allocates memory which is attached to the mm_struct and freed
284             // when the mm_struct is freed.
285             ret = __mmu_notifier_register(NULL, current->mm);
286             if (ret)
287                 return errno_to_nv_status(ret);
288         #else
289             UVM_ASSERT(0);
290         #endif
291     }
292 
293     return NV_OK;
294 }
295 
296 void uvm_va_space_mm_unregister(uvm_va_space_t *va_space)
297 {
298     uvm_va_space_mm_t *va_space_mm = &va_space->va_space_mm;
299 
300     // We can't hold the VA space lock or mmap_lock because
301     // uvm_va_space_mm_shutdown() waits for retainers which may take
302     // these locks.
303     uvm_assert_unlocked_order(UVM_LOCK_ORDER_MMAP_LOCK);
304     uvm_assert_unlocked_order(UVM_LOCK_ORDER_VA_SPACE);
305 
306     uvm_va_space_mm_shutdown(va_space);
307     UVM_ASSERT(va_space_mm->retained_count == 0);
308 
309     // Only happens if uvm_va_space_mm_register() fails
310     if (!va_space_mm->mm)
311         return;
312 
313     // At this point the mm is still valid because uvm_mm_release()
314     // hasn't yet called mmput(). uvm_hmm_va_space_destroy() will kill
315     // all the va_blocks along with any associated gpu_chunks, so we
316     // need to make sure these chunks are free. However freeing them
317     // requires a valid mm so we can call migrate_vma_setup(), so we
318     // do that here.
319     // TODO: Bug 3902536: [UVM-HMM] add code to migrate GPU memory
320     // without having a va_block
321     if (uvm_hmm_is_enabled(va_space))
322         uvm_hmm_evict_va_blocks(va_space);
323 
324     if (uvm_va_space_mm_enabled(va_space)) {
325         if (UVM_ATS_IBM_SUPPORTED_IN_DRIVER() && g_uvm_global.ats.enabled)
326             uvm_mmu_notifier_unregister(va_space_mm);
327         uvm_mmdrop(va_space_mm->mm);
328     }
329 }
330 
331 struct mm_struct *uvm_va_space_mm_retain(uvm_va_space_t *va_space)
332 {
333     uvm_va_space_mm_t *va_space_mm = &va_space->va_space_mm;
334     struct mm_struct *mm = NULL;
335 
336     if (!uvm_va_space_mm_enabled(va_space))
337         return NULL;
338 
339     uvm_spin_lock(&va_space_mm->lock);
340 
341     if (!uvm_va_space_mm_alive(va_space_mm))
342         goto out;
343 
344     ++va_space_mm->retained_count;
345 
346     mm = va_space_mm->mm;
347     UVM_ASSERT(mm);
348 
349 out:
350 
351     // uvm_api_mm_init() holds a reference
352     if (mm)
353         UVM_ASSERT(atomic_read(&mm->mm_users) > 0);
354 
355     uvm_spin_unlock(&va_space_mm->lock);
356 
357     return mm;
358 }
359 
360 struct mm_struct *uvm_va_space_mm_or_current_retain(uvm_va_space_t *va_space)
361 {
362     // We should only attempt to use current->mm from a user thread
363     UVM_ASSERT(!(current->flags & PF_KTHREAD));
364 
365     // current->mm is NULL when we're in process teardown. In that case it
366     // doesn't make sense to use any mm.
367     if (!current->mm)
368         return NULL;
369 
370     // If !uvm_va_space_mm_enabled() we use current->mm on the ioctl
371     // paths. In that case we don't need to mmget(current->mm) because
372     // the current thread mm is always valid. On
373     // uvm_va_space_mm_enabled() systems we skip trying to retain the
374     // mm if it is current->mm because userspace may not have
375     // initialised the mm fd but UVM callers on the ioctl path still
376     // assume retaining current->mm will succeed.
377     if (!uvm_va_space_mm_enabled(va_space))
378         return current->mm;
379 
380     return uvm_va_space_mm_retain(va_space);
381 }
382 
383 void uvm_va_space_mm_release(uvm_va_space_t *va_space)
384 {
385     uvm_va_space_mm_t *va_space_mm = &va_space->va_space_mm;
386 
387     UVM_ASSERT(uvm_va_space_mm_enabled(va_space));
388 
389     // The mm must not have been torn down while we have it retained
390     UVM_ASSERT(va_space_mm->mm);
391 
392     uvm_spin_lock(&va_space_mm->lock);
393 
394     UVM_ASSERT(va_space_mm->retained_count > 0);
395     --va_space_mm->retained_count;
396 
397     // If we're the last retainer on a dead mm, signal any potential waiters
398     if (va_space_mm->retained_count == 0 && !uvm_va_space_mm_alive(va_space_mm)) {
399         uvm_spin_unlock(&va_space_mm->lock);
400 
401         // There could be a thread in uvm_va_space_mm_shutdown()
402         // waiting on us, so wake it up.
403         wake_up(&va_space_mm->last_retainer_wait_queue);
404     }
405     else {
406         uvm_spin_unlock(&va_space_mm->lock);
407     }
408 }
409 
410 void uvm_va_space_mm_or_current_release(uvm_va_space_t *va_space, struct mm_struct *mm)
411 {
412     if (!uvm_va_space_mm_enabled(va_space) || !mm)
413         return;
414 
415     uvm_va_space_mm_release(va_space);
416 }
417 
418 static void uvm_va_space_mm_shutdown(uvm_va_space_t *va_space)
419 {
420     uvm_va_space_mm_t *va_space_mm = &va_space->va_space_mm;
421     uvm_gpu_va_space_t *gpu_va_space;
422     uvm_gpu_t *gpu;
423     uvm_global_processor_mask_t gpus_to_flush;
424     LIST_HEAD(deferred_free_list);
425 
426     uvm_va_space_down_write(va_space);
427 
428     // Prevent future registrations of any kind. We'll be iterating over all
429     // GPUs and GPU VA spaces below but taking and dropping the VA space lock.
430     // It's ok for other threads to unregister those objects, but not to
431     // register new ones.
432     //
433     // We also need to prevent new channel work from arriving since we're trying
434     // to stop memory accesses.
435     va_space->disallow_new_registers = true;
436 
437     uvm_va_space_downgrade_write_rm(va_space);
438 
439     // Stop channels to prevent new accesses and new faults on non-MPS
440     uvm_va_space_stop_all_user_channels(va_space);
441 
442     uvm_va_space_up_read_rm(va_space);
443 
444     // Detach all channels to prevent pending untranslated faults from getting
445     // to this VA space. This also removes those channels from the VA space and
446     // puts them on the deferred free list, so only one thread will do this.
447     uvm_va_space_down_write(va_space);
448     uvm_va_space_detach_all_user_channels(va_space, &deferred_free_list);
449     uvm_va_space_global_gpus_in_mask(va_space, &gpus_to_flush, &va_space->faultable_processors);
450     uvm_global_mask_retain(&gpus_to_flush);
451     uvm_va_space_up_write(va_space);
452 
453     // Flush the fault buffer on all GPUs. This will avoid spurious
454     // cancels of stale pending translated faults after we set
455     // UVM_VA_SPACE_MM_STATE_RELEASED later.
456     for_each_global_gpu_in_mask(gpu, &gpus_to_flush)
457         uvm_gpu_fault_buffer_flush(gpu);
458 
459     uvm_global_mask_release(&gpus_to_flush);
460 
461     // Call nvUvmInterfaceUnsetPageDirectory. This has no effect on non-MPS.
462     // Under MPS this guarantees that no new GPU accesses will be made using
463     // this mm.
464     //
465     // We need only one thread to make this call, but two threads in here could
466     // race for it, or we could have one thread in here and one in
467     // destroy_gpu_va_space. Serialize these by starting in write mode then
468     // downgrading to read.
469     uvm_va_space_down_write(va_space);
470     uvm_va_space_downgrade_write_rm(va_space);
471     for_each_gpu_va_space(gpu_va_space, va_space)
472         uvm_gpu_va_space_unset_page_dir(gpu_va_space);
473     uvm_va_space_up_read_rm(va_space);
474 
475     // The above call to uvm_gpu_va_space_unset_page_dir handles the GPU VA
476     // spaces which are known to be registered. However, we could've raced with
477     // a concurrent uvm_va_space_unregister_gpu_va_space, giving this sequence:
478     //
479     // unregister_gpu_va_space                  uvm_va_space_mm_shutdown
480     //     uvm_va_space_down_write
481     //     remove_gpu_va_space
482     //     uvm_va_space_up_write
483     //                                          uvm_va_space_down_write(va_space);
484     //                                          // No GPU VA spaces
485     //                                          Unlock, return
486     //     uvm_deferred_free_object_list
487     //         uvm_gpu_va_space_unset_page_dir
488     //
489     // We have to be sure that all accesses in this GPU VA space are done before
490     // returning, so we have to wait for the other thread to finish its
491     // uvm_gpu_va_space_unset_page_dir call.
492     //
493     // We can be sure that num_pending will eventually go to zero because we've
494     // prevented new GPU VA spaces from being registered above.
495     wait_event(va_space->gpu_va_space_deferred_free.wait_queue,
496                atomic_read(&va_space->gpu_va_space_deferred_free.num_pending) == 0);
497 
498     // Now that there won't be any new GPU faults, prevent subsequent retainers
499     // from accessing this mm.
500     uvm_spin_lock(&va_space_mm->lock);
501     va_space_mm->state = UVM_VA_SPACE_MM_STATE_RELEASED;
502     uvm_spin_unlock(&va_space_mm->lock);
503 
504     // Finish channel destroy. This can be done at any point after detach as
505     // long as we don't hold the VA space lock.
506     uvm_deferred_free_object_list(&deferred_free_list);
507 
508     // Flush out all pending retainers
509     wait_event(va_space_mm->last_retainer_wait_queue, va_space_mm->retained_count == 0);
510 }
511 
512 static NV_STATUS mm_read64(struct mm_struct *mm, NvU64 addr, NvU64 *val)
513 {
514     long ret;
515     struct page *page;
516     NvU64 *mapping;
517 
518     UVM_ASSERT(IS_ALIGNED(addr, sizeof(*val)));
519 
520     uvm_down_read_mmap_lock(mm);
521     ret = NV_PIN_USER_PAGES_REMOTE(mm, (unsigned long)addr, 1, 0, &page, NULL, NULL);
522     uvm_up_read_mmap_lock(mm);
523 
524     if (ret < 0)
525         return errno_to_nv_status(ret);
526 
527     UVM_ASSERT(ret == 1);
528 
529     mapping = (NvU64 *)((char *)kmap(page) + (addr % PAGE_SIZE));
530     *val = *mapping;
531     kunmap(page);
532     NV_UNPIN_USER_PAGE(page);
533 
534     return NV_OK;
535 }
536 
537 NV_STATUS uvm_test_va_space_mm_retain(UVM_TEST_VA_SPACE_MM_RETAIN_PARAMS *params, struct file *filp)
538 {
539     uvm_va_space_t *va_space = NULL;
540     struct mm_struct *mm = NULL;
541     NV_STATUS status = NV_OK;
542 
543     if (!IS_ALIGNED(params->addr, sizeof(params->val_before)))
544         return NV_ERR_INVALID_ARGUMENT;
545 
546     uvm_mutex_lock(&g_uvm_global.va_spaces.lock);
547 
548     list_for_each_entry(va_space, &g_uvm_global.va_spaces.list, list_node) {
549         if ((uintptr_t)va_space == params->va_space_ptr) {
550             mm = uvm_va_space_mm_retain(va_space);
551             break;
552         }
553     }
554 
555     uvm_mutex_unlock(&g_uvm_global.va_spaces.lock);
556 
557     if ((uintptr_t)va_space != params->va_space_ptr)
558         return NV_ERR_MISSING_TABLE_ENTRY;
559 
560     if (!mm)
561         return NV_ERR_PAGE_TABLE_NOT_AVAIL;
562 
563     status = mm_read64(mm, params->addr, &params->val_before);
564 
565     if (status == NV_OK && params->sleep_us) {
566         usleep_range(params->sleep_us, params->sleep_us + 1000);
567         status = mm_read64(mm, params->addr, &params->val_after);
568     }
569 
570     uvm_va_space_mm_release(va_space);
571     return status;
572 }
573 
574 NV_STATUS uvm_test_va_space_mm_or_current_retain(UVM_TEST_VA_SPACE_MM_OR_CURRENT_RETAIN_PARAMS *params,
575                                                  struct file *filp)
576 {
577     uvm_va_space_t *va_space = uvm_va_space_get(filp);
578     struct mm_struct *mm;
579     NV_STATUS status = NV_OK;
580 
581     mm = uvm_va_space_mm_or_current_retain(va_space);
582     if (!mm)
583         return NV_ERR_PAGE_TABLE_NOT_AVAIL;
584 
585     if (params->retain_done_ptr) {
586         NvU64 flag = true;
587 
588         if (nv_copy_to_user((void __user *)params->retain_done_ptr, &flag, sizeof(flag)))
589             status = NV_ERR_INVALID_ARGUMENT;
590     }
591 
592     if (status == NV_OK && params->sleep_us)
593             usleep_range(params->sleep_us, params->sleep_us + 1000);
594 
595     uvm_va_space_mm_or_current_release(va_space, mm);
596 
597     return status;
598 }
599