From 3b6f3e2eebaa48412af64620a0f157ee19de0070 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Fri, 2 Feb 2018 15:06:17 -0500
Subject: [PATCH] expose UVM API

---
 src/memory.jl | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/memory.jl b/src/memory.jl
index 82c546b..3e8e998 100644
--- a/src/memory.jl
+++ b/src/memory.jl
@@ -5,7 +5,7 @@ export Mem
 module Mem
 
 using CUDAdrv
-import CUDAdrv: @apicall, CuStream_t
+import CUDAdrv: @apicall, CuStream_t, CuDevice_t
 
 
 
@@ -143,19 +143,47 @@ function transfer end
 
 ## pointer-based
 
+@enum(CUmem_attach, ATTACH_GLOBAL = 0x01,
+                    ATTACH_HOST   = 0x02)
+                    #ATTACH_SINGLE = 0x04) # Defined but not valid
+
 """
     alloc(bytes::Integer)
 
 Allocate `bytesize` bytes of memory.
 """
-function alloc(bytesize::Integer)
+function alloc(bytesize::Integer, managed=false; flags::CUmem_attach=ATTACH_GLOBAL)
     bytesize == 0 && throw(ArgumentError("invalid amount of memory requested"))
 
     ptr_ref = Ref{Ptr{Cvoid}}()
-    @apicall(:cuMemAlloc, (Ptr{Ptr{Cvoid}}, Csize_t), ptr_ref, bytesize)
+    if !managed
+        @apicall(:cuMemAlloc, (Ptr{Ptr{Cvoid}}, Csize_t), ptr_ref, bytesize)
+    else
+        @apicall(:cuMemAllocManaged, (Ptr{Ptr{Cvoid}}, Csize_t, Cuint), ptr_ref, bytesize, flags)
+    end
     return Buffer(ptr_ref[], bytesize, CuCurrentContext())
 end
 
+function prefetch(buf::Buffer, bytes=buf.bytesize; stream::CuStream=CuDefaultStream())
+    bytes > buf.bytesize && throw(BoundsError(buf, bytes))
+    dev = device(buf.ctx)
+    @apicall(:cuMemPrefetchAsync, (Ptr{Cvoid}, Csize_t, CuDevice_t, CuStream_t),
+             buf, bytes, dev, stream)
+end
+
+@enum(CUmem_advise, ADVISE_SET_READ_MOSTLY          = 0x01,
+                    ADVISE_UNSET_READ_MOSTLY        = 0x02,
+                    ADVISE_SET_PREFERRED_LOCATION   = 0x03,
+                    ADVISE_UNSET_PREFERRED_LOCATION = 0x04,
+                    ADVISE_SET_ACCESSED_BY          = 0x05,
+                    ADVISE_UNSET_ACCESSED_BY        = 0x06)
+
+function advise(buf::Buffer, advice::CUmem_advise, bytes=buf.bytesize, device=device(buf.ctx))
+    bytes > buf.bytesize && throw(BoundsError(buf, bytes))
+    @apicall(:cuMemAdvise, (Ptr{Cvoid}, Csize_t, Cuint, CuDevice_t),
+             buf, bytes, advice, device)
+end
+
 function free(buf::Buffer)
     @apicall(:cuMemFree, (Ptr{Cvoid},), buf.ptr)
     return