Memory improvements (#395)

* Add an OutOfGPUMemoryError type with additional info.
* Relax allocation limits.

Author: maleadt

lib/level-zero/memory.jl

@@ -44,14 +44,24 @@ struct DeviceBuffer <: AbstractBuffer
     device::ZeDevice
 end
 
-function device_alloc(ctx::ZeContext, dev::ZeDevice, bytesize::Integer, alignment::Integer=1;
-                      flags=0, ordinal::Integer=0)
-    desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
-
-    ptr_ref = Ref{Ptr{Cvoid}}()
-    zeMemAllocDevice(ctx, desc_ref, bytesize, alignment, dev, ptr_ref)
-
-    return DeviceBuffer(reinterpret(ZePtr{Cvoid}, ptr_ref[]), bytesize, ctx, dev)
+function device_alloc(ctx::ZeContext, dev::ZeDevice, bytesize::Integer,
+                      alignment::Integer=1; flags=0, ordinal::Integer=0)
+    relaxed_allocation_ref = Ref(ze_relaxed_allocation_limits_exp_desc_t(;
+        flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE,
+    ))
+    GC.@preserve relaxed_allocation_ref begin
+        desc_ref = if bytesize > properties(dev).maxMemAllocSize
+            pNext = Base.unsafe_convert(Ptr{Cvoid}, relaxed_allocation_ref)
+            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal, pNext))
+        else
+            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
+        end
+
+        ptr_ref = Ref{Ptr{Cvoid}}()
+        zeMemAllocDevice(ctx, desc_ref, bytesize, alignment, dev, ptr_ref)
+
+        return DeviceBuffer(reinterpret(ZePtr{Cvoid}, ptr_ref[]), bytesize, ctx, dev)
+    end
 end
 
 Base.pointer(buf::DeviceBuffer) = buf.ptr
@@ -84,7 +94,8 @@ struct HostBuffer <: AbstractBuffer
     context::ZeContext
 end
 
-function host_alloc(ctx::ZeContext, bytesize::Integer, alignment::Integer=1; flags=0)
+function host_alloc(ctx::ZeContext, bytesize::Integer, alignment::Integer=1;
+                    flags=0)
     desc_ref = Ref(ze_host_mem_alloc_desc_t(; flags))
 
     ptr_ref = Ref{Ptr{Cvoid}}()
@@ -122,17 +133,27 @@ struct SharedBuffer <: AbstractBuffer
     device::Union{Nothing,ZeDevice}
 end
 
-function shared_alloc(ctx::ZeContext, dev::Union{Nothing,ZeDevice}, bytesize::Integer,
-                      alignment::Integer=1; host_flags=0,
+function shared_alloc(ctx::ZeContext, dev::Union{Nothing,ZeDevice},
+                      bytesize::Integer, alignment::Integer=1; host_flags=0,
                       device_flags=0, ordinal::Integer=0)
-    device_desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags=device_flags, ordinal))
-    host_desc_ref = Ref(ze_host_mem_alloc_desc_t(; flags=host_flags))
-
-    ptr_ref = Ref{Ptr{Cvoid}}()
-    zeMemAllocShared(ctx, device_desc_ref, host_desc_ref, bytesize, alignment,
-                     something(dev, C_NULL), ptr_ref)
-
-    return SharedBuffer(reinterpret(ZePtr{Cvoid}, ptr_ref[]), bytesize, ctx, dev)
+    relaxed_allocation_ref = Ref(ze_relaxed_allocation_limits_exp_desc_t(;
+        flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE,
+    ))
+    GC.@preserve relaxed_allocation_ref begin
+        device_desc_ref = if dev !== nothing && bytesize > properties(dev).maxMemAllocSize
+            pNext = Base.unsafe_convert(Ptr{Cvoid}, relaxed_allocation_ref)
+            Ref(ze_device_mem_alloc_desc_t(; flags=device_flags, ordinal, pNext))
+        else
+            Ref(ze_device_mem_alloc_desc_t(; flags=device_flags, ordinal))
+        end
+        host_desc_ref = Ref(ze_host_mem_alloc_desc_t(; flags=host_flags))
+
+        ptr_ref = Ref{Ptr{Cvoid}}()
+        zeMemAllocShared(ctx, device_desc_ref, host_desc_ref, bytesize, alignment,
+                        something(dev, C_NULL), ptr_ref)
+
+        return SharedBuffer(reinterpret(ZePtr{Cvoid}, ptr_ref[]), bytesize, ctx, dev)
+    end
 end
 
 Base.pointer(buf::SharedBuffer) = buf.ptr

lib/level-zero/oneL0.jl

@@ -24,6 +24,7 @@ include("libze.jl")
 #
 # TODO: is it really required to (1) zero-initialize memory and (2) set the stype field?
 #       none of these are documented...
+# TODO: add support for conveniently linking extension objects in pNext
 for (structure_type_enum, _) in CEnum.name_value_pairs(ze_structure_type_t)
     structure_type_name = string(structure_type_enum)
     @assert startswith(structure_type_name, "ZE_STRUCTURE_TYPE")

src/pool.jl

@@ -1,3 +1,30 @@
+"""
+    OutOfGPUMemoryError()
+
+An operation allocated too much GPU memory.
+"""
+struct OutOfGPUMemoryError <: Exception
+  sz::Int
+  dev::ZeDevice
+
+  function OutOfGPUMemoryError(sz::Integer=0, dev::ZeDevice=device())
+    new(sz, dev)
+  end
+end
+
+function Base.showerror(io::IO, err::OutOfGPUMemoryError)
+    print(io, "Out of GPU memory")
+    if err.sz > 0
+      print(io, " trying to allocate $(Base.format_bytes(err.sz))")
+    end
+    print(" on device $(properties(err.dev).name)")
+    if length(memory_properties(err.dev)) == 1
+        # XXX: how to handle multiple memories?
+        print(" with $(Base.format_bytes(only(memory_properties(err.dev)).totalSize))")
+    end
+    return io
+end
+
 function allocate(::Type{oneL0.DeviceBuffer}, ctx, dev, bytes::Int, alignment::Int)
     bytes == 0 && return oneL0.DeviceBuffer(ZE_NULL, bytes, ctx, dev)