1 files changed, 6 insertions, 7 deletions

diff --git a/drivers/gpu/nvgpu/os/linux/ioctl_as.c b/drivers/gpu/nvgpu/os/linux/ioctl_as.c
index af6cdb5b..6348bb2a 100644
--- a/drivers/gpu/nvgpu/os/linux/ioctl_as.c
+++ b/drivers/gpu/nvgpu/os/linux/ioctl_as.c
@@ -365,9 +365,9 @@ static int nvgpu_as_ioctl_write_swap_buffer(
         // (Assuming that NVMe DRAM acceses are uncached)
         gk20a_mm_l2_flush(g, false);
 
-        // Copy out (blocking) TODO: non-blocking
+        // Copy out (non-blocking)
         // Could fail on inaccessible swap device, etc
-        err = copy_out(m->os_priv.sgt);
+        err = copy_out(m->os_priv.sgt, m);
 
 out:
         return err;
@@ -393,7 +393,7 @@ static int nvgpu_as_ioctl_write_swap_buffer_finish(
         nvgpu_log_fn(g, " ");
 
         // Wait for the pages to get written out
-        //wait_for_completion_io(m->os_priv.swap_completion);
+        wait_for_completion_io(&m->os_priv.swap_io_done);
 
         // Unpin needs to happen after copy out is done
         // (No return value check as it's a void function)
@@ -448,9 +448,8 @@ static int nvgpu_as_ioctl_read_swap_buffer(
         // Do any bookeeping not done by gk20a_mm_pin()
         m->os_priv.sgt = sgt;
 
-        // Reload page contents from disk (blocking)
-        // TODO: non-blocking
-        err = copy_in(sgt);
+        // Reload page contents from disk (non-blocking)
+        err = copy_in(sgt, m);
         if (err) {
                 int err2;
                 // Rollback pinning and allocation
@@ -487,7 +486,7 @@ static int nvgpu_as_ioctl_read_swap_buffer_finish(
         // Invalidate L2 so that TLB refill does not load stale PT
         gk20a_mm_l2_flush(g, true);
         // Wait for read to complete if it hasn't yet
-        //wait_for_completion_io(m->os_priv.swap_completion);
+        wait_for_completion_io(&m->os_priv.swap_io_done);
 
         return err;
 }