Module: TensorRT::CUDA

Extended by:
FFI::Library
Defined in:
lib/tensorrt/cuda.rb

Constant Summary collapse

CUDA_LIBS =
%w[
  libcudart.so
  libcudart.so.12
  libcudart.so.11
].freeze
MEMCPY_HOST_TO_DEVICE =
1
MEMCPY_DEVICE_TO_HOST =
2
Error =
Class.new(StandardError)

Class Method Summary collapse

Class Method Details

.free(ptr) ⇒ Object

Raises:



43
44
45
46
47
# File 'lib/tensorrt/cuda.rb', line 43

def free(ptr)
  err = cudaFree(FFI::Pointer.new(ptr))

  raise Error, "cudaFree failed with error #{err}" unless err.zero?
end

.malloc(size) ⇒ Object

Raises:



34
35
36
37
38
39
40
41
# File 'lib/tensorrt/cuda.rb', line 34

def malloc(size)
  ptr = FFI::MemoryPointer.new(:pointer)
  err = cudaMalloc(ptr, size)

  raise Error, "cudaMalloc failed with error #{err}" unless err.zero?

  ptr.read_pointer.address
end

.memcpy_dtoh(dst_ptr, src, size) ⇒ Object

Raises:



55
56
57
58
59
# File 'lib/tensorrt/cuda.rb', line 55

def memcpy_dtoh(dst_ptr, src, size)
  err = cudaMemcpy(dst_ptr, FFI::Pointer.new(src), size, MEMCPY_DEVICE_TO_HOST)

  raise Error, "cudaMemcpy D2H failed with error #{err}" unless err.zero?
end

.memcpy_dtoh_async(dst_ptr, src, size, stream) ⇒ Object

Raises:



67
68
69
70
71
# File 'lib/tensorrt/cuda.rb', line 67

def memcpy_dtoh_async(dst_ptr, src, size, stream)
  err = cudaMemcpyAsync(dst_ptr, FFI::Pointer.new(src), size, MEMCPY_DEVICE_TO_HOST, FFI::Pointer.new(stream))

  raise Error, "cudaMemcpyAsync D2H failed with error #{err}" unless err.zero?
end

.memcpy_htod(dst, src_ptr, size) ⇒ Object

Raises:



49
50
51
52
53
# File 'lib/tensorrt/cuda.rb', line 49

def memcpy_htod(dst, src_ptr, size)
  err = cudaMemcpy(FFI::Pointer.new(dst), src_ptr, size, MEMCPY_HOST_TO_DEVICE)

  raise Error, "cudaMemcpy H2D failed with error #{err}" unless err.zero?
end

.memcpy_htod_async(dst, src_ptr, size, stream) ⇒ Object

Raises:



61
62
63
64
65
# File 'lib/tensorrt/cuda.rb', line 61

def memcpy_htod_async(dst, src_ptr, size, stream)
  err = cudaMemcpyAsync(FFI::Pointer.new(dst), src_ptr, size, MEMCPY_HOST_TO_DEVICE, FFI::Pointer.new(stream))

  raise Error, "cudaMemcpyAsync H2D failed with error #{err}" unless err.zero?
end

.stream_synchronize(stream) ⇒ Object

Raises:



79
80
81
82
83
# File 'lib/tensorrt/cuda.rb', line 79

def stream_synchronize(stream)
  err = cudaStreamSynchronize(FFI::Pointer.new(stream))

  raise Error, "cudaStreamSynchronize failed with error #{err}" unless err.zero?
end

.synchronizeObject

Raises:



73
74
75
76
77
# File 'lib/tensorrt/cuda.rb', line 73

def synchronize
  err = cudaDeviceSynchronize

  raise Error, "cudaDeviceSynchronize failed with error #{err}" unless err.zero?
end