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
- .free(ptr) ⇒ Object
- .malloc(size) ⇒ Object
- .memcpy_dtoh(dst_ptr, src, size) ⇒ Object
- .memcpy_dtoh_async(dst_ptr, src, size, stream) ⇒ Object
- .memcpy_htod(dst, src_ptr, size) ⇒ Object
- .memcpy_htod_async(dst, src_ptr, size, stream) ⇒ Object
- .stream_synchronize(stream) ⇒ Object
- .synchronize ⇒ Object
Class Method Details
.free(ptr) ⇒ Object
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
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
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
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
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
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
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 |
.synchronize ⇒ Object
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 |