Ensure essential function calls don't allocate.

JuliaGPU · Apr 25, 2024 · 51fb828 · 51fb828
1 parent a534c10
commit 51fb828
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 12 deletions.
diff --git a/lib/cudadrv/context.jl b/lib/cudadrv/context.jl
@@ -56,7 +56,7 @@ has no context bound to it, or if the bound context has been destroyed.
 """
 global function current_context()
  handle_ref = Ref{CUcontext}()
- cuCtxGetCurrent(handle_ref)
+ @inline cuCtxGetCurrent(handle_ref) # JuliaGPU/CUDA.jl#2347
  handle_ref[] == C_NULL && throw(UndefRefError())
  CuContext(handle_ref[])
 end

diff --git a/lib/cudadrv/libcuda.jl b/lib/cudadrv/libcuda.jl
@@ -31,7 +31,7 @@ end
  end
 end
 
-function check(f)
+@inline function check(f)
  res = f()
  if res != SUCCESS
  throw_api_error(res)

diff --git a/lib/nvml/libnvml.jl b/lib/nvml/libnvml.jl
@@ -23,7 +23,7 @@ function initialize_context()
  end
 end
 
-function check(f)
+@inline function check(f)
  res = f()
  if res != NVML_SUCCESS
  throw_api_error(res)

diff --git a/lib/utils/call.jl b/lib/utils/call.jl
@@ -196,19 +196,15 @@ function ccall_macro_lower(func, rettype, types, args, nreq)
  for (typ, arg) in zip(types, args)
  var = gensym("$(func)_cconvert")
  push!(cconvert_args, var)
- push!(cconvert_exprs, quote
- $var = Base.cconvert($(esc(typ)), $(esc(arg)))
- end)
+ push!(cconvert_exprs, :($var = Base.cconvert($(esc(typ)), $(esc(arg)))))
  end
 
  unsafe_convert_exprs = []
  unsafe_convert_args = []
  for (typ, arg) in zip(types, cconvert_args)
  var = gensym("$(func)_unsafe_convert")
  push!(unsafe_convert_args, var)
- push!(unsafe_convert_exprs, quote
- $var = Base.unsafe_convert($(esc(typ)), $arg)
- end)
+ push!(unsafe_convert_exprs, :($var = Base.unsafe_convert($(esc(typ)), $arg)))
  end
 
  call = quote
@@ -222,9 +218,12 @@ function ccall_macro_lower(func, rettype, types, args, nreq)
  end
 
  quote
+ # the added operations push our simple ccall wrappers over the inlining limit
+ @inline
+
  $(cconvert_exprs...)
 
- GC.@preserve $(cconvert_args...)  $(call)
+ GC.@preserve $(cconvert_args...) $(call)
  end
 end
 

diff --git a/res/wrap/libcuda_prologue.jl b/res/wrap/libcuda_prologue.jl
@@ -26,7 +26,7 @@ end
  end
 end
 
-function check(f)
+@inline function check(f)
  res = f()
  if res != SUCCESS
  throw_api_error(res)

diff --git a/res/wrap/libnvml_prologue.jl b/res/wrap/libnvml_prologue.jl
@@ -18,7 +18,7 @@ function initialize_context()
  end
 end
 
-function check(f)
+@inline function check(f)
  res = f()
  if res != NVML_SUCCESS
  throw_api_error(res)

diff --git a/test/core/initialization.jl b/test/core/initialization.jl
@@ -177,3 +177,11 @@ end
  proc, out, err = julia_exec(`-e $script`, "CUDA_VISIBLE_DEVICES"=>"-1")
  @test success(proc)
 end
+
+
+## allocations
+
+@test @allocated(current_context()) == 0
+@test @allocated(context()) == 0
+@test @allocated(stream()) == 0
+@test @allocated(device()) == 0