tensorrt 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 165592be0806398f6bb72846feb414eae2947e592cdba7876f91a92a1850d352
+  data.tar.gz: 7f585077dec7ee5de44362d6b93a365ec696b112cd63085006ea6c613e392870
+SHA512:
+  metadata.gz: 33fff8454d89035f4785c556e13207b132ae8677839b3eee073cc69bad3ab621390b41e96fb4eb94542ec00ca29f7c67fee04e24597e971d2f4244af3060d895
+  data.tar.gz: 45a0cf35ed380cb77a3d944f8355cd8945184b3584f137d5e6de8f0d1cf32f34778a53f2c782253e7f4e91310fdf0e8134a58db0f8cfc469670212d0890c3acb

data/README.md

@@ -0,0 +1,182 @@
+# tensorrt-rb
+
+Minimal TensorRT bindings for Ruby using Rice (C++ bindings).
+
+## Requirements
+
+- Linux (x86_64 or aarch64)
+- Ruby >= 3.0
+- TensorRT (with headers and libraries)
+- CUDA runtime
+- Rice gem (`gem install rice`)
+
+## Installation
+
+```bash
+cd tensorrt-rb
+
+# Set paths if not in standard locations
+export TENSORRT_INCLUDE=/path/to/tensorrt/include
+export TENSORRT_LIB=/path/to/tensorrt/lib
+export CUDA_INCLUDE=/usr/local/cuda/include
+export CUDA_LIB=/usr/local/cuda/lib64
+
+# Build
+rake compile
+
+# Or install as gem
+gem build tensorrt.gemspec
+gem install tensorrt-*.gem
+```
+
+### Default Library Paths
+
+**x86_64:**
+- TensorRT: `/usr/include/x86_64-linux-gnu`, `/usr/lib/x86_64-linux-gnu`
+- CUDA: `/usr/local/cuda/include`, `/usr/local/cuda/lib64`
+
+**aarch64:**
+- TensorRT: `/usr/include/aarch64-linux-gnu`, `/usr/lib/aarch64-linux-gnu`
+- CUDA: `/usr/local/cuda/include`, `/usr/local/cuda/lib64`
+
+## API
+
+### TensorRT::Engine
+
+```ruby
+engine = TensorRT::Engine.new(path, verbose: false)
+
+# Tensor info
+engine.num_io_tensors                  # Number of input/output tensors
+engine.get_tensor_name(index)          # Tensor name by index
+engine.is_input?(name)                 # Check if tensor is input
+engine.get_tensor_shape(name)          # Shape as array [1, 3, 640, 640]
+engine.get_tensor_bytes(name)          # Size in bytes
+
+# Memory binding
+engine.set_tensor_address(name, device_ptr)
+
+# Inference
+engine.execute                         # Synchronous (blocking)
+engine.enqueue                         # Asynchronous (non-blocking)
+
+# Stream management
+engine.get_stream                      # CUDA stream handle (uint64)
+engine.stream_synchronize              # Wait for stream completion
+```
+
+### TensorRT::CUDA
+
+```ruby
+# Memory allocation
+ptr = TensorRT::CUDA.malloc(bytes)
+TensorRT::CUDA.free(ptr)
+
+# Synchronous copy
+TensorRT::CUDA.memcpy_htod(device_ptr, host_ptr, bytes)  # Host → Device
+TensorRT::CUDA.memcpy_dtoh(host_ptr, device_ptr, bytes)  # Device → Host
+
+# Asynchronous copy
+TensorRT::CUDA.memcpy_htod_async(device_ptr, host_ptr, bytes, stream)
+TensorRT::CUDA.memcpy_dtoh_async(host_ptr, device_ptr, bytes, stream)
+
+# Synchronization
+TensorRT::CUDA.synchronize                    # All operations
+TensorRT::CUDA.stream_synchronize(stream)     # Specific stream
+```
+
+## Examples
+
+### Synchronous Inference
+
+```ruby
+require "tensorrt"
+
+engine = TensorRT::Engine.new("model.engine")
+
+# Allocate GPU memory
+input_bytes = engine.get_tensor_bytes("input")
+output_bytes = engine.get_tensor_bytes("output")
+output_size = engine.get_tensor_shape("output").reduce(1, :*)
+
+input_device = TensorRT::CUDA.malloc(input_bytes)
+output_device = TensorRT::CUDA.malloc(output_bytes)
+engine.set_tensor_address("input", input_device)
+engine.set_tensor_address("output", output_device)
+
+# Prepare input data
+input_data = preprocess_image(image_path)  # Returns Numo::SFloat
+input_host = FFI::MemoryPointer.new(:float, input_data.size)
+input_host.write_bytes(input_data.to_binary)
+
+# Copy input to GPU
+TensorRT::CUDA.memcpy_htod(input_device, input_host, input_bytes)
+
+# Run inference
+engine.execute
+
+# Copy output from GPU
+output_host = FFI::MemoryPointer.new(:float, output_size)
+TensorRT::CUDA.memcpy_dtoh(output_host, output_device, output_bytes)
+output_data = output_host.read_array_of_float(output_size)
+
+# Cleanup
+TensorRT::CUDA.free(input_device)
+TensorRT::CUDA.free(output_device)
+```
+
+### Pipelined Async Inference
+
+Overlap CPU preprocessing with GPU inference for maximum throughput:
+
+```ruby
+require "tensorrt"
+
+engine = TensorRT::Engine.new("model.engine")
+stream = engine.get_stream
+
+# Allocate GPU memory
+input_bytes = engine.get_tensor_bytes("input")
+output_bytes = engine.get_tensor_bytes("output")
+output_size = engine.get_tensor_shape("output").reduce(1, :*)
+
+input_device = TensorRT::CUDA.malloc(input_bytes)
+output_device = TensorRT::CUDA.malloc(output_bytes)
+engine.set_tensor_address("input", input_device)
+engine.set_tensor_address("output", output_device)
+
+# Allocate host buffers
+input_host = FFI::MemoryPointer.new(:float, input_bytes / 4)
+output_host = FFI::MemoryPointer.new(:float, output_size)
+
+# Preload first image
+current_image = preprocess_image(images[0])
+
+images.each_with_index do |image_path, i|
+  # Copy current image to GPU (async)
+  input_host.write_bytes(current_image.to_binary)
+  TensorRT::CUDA.memcpy_htod_async(input_device, input_host, input_bytes, stream)
+
+  # Start async inference
+  engine.enqueue
+
+  # Preprocess next image on CPU while GPU is busy
+  next_image = preprocess_image(images[i + 1]) if i < images.size - 1
+
+  # Wait for GPU inference to complete
+  engine.stream_synchronize
+
+  # Copy output from GPU
+  TensorRT::CUDA.memcpy_dtoh(output_host, output_device, output_bytes)
+  output_data = output_host.read_array_of_float(output_size)
+
+  # Process results
+  process_detections(output_data)
+
+  current_image = next_image
+end
+
+# Cleanup
+TensorRT::CUDA.free(input_device)
+TensorRT::CUDA.free(output_device)
+```

data/Rakefile

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "rake/extensiontask"
+
+task default: :compile
+
+Rake::ExtensionTask.new("tensorrt_rb") do |ext|
+  ext.lib_dir = "lib/tensorrt_rb"
+end
+
+desc "Build and install the gem"
+task :install => :compile do
+  sh "gem build tensorrt.gemspec && gem install tensorrt-*.gem"
+end
+
+desc "Clean build artifacts"
+task :clean do
+  rm_rf "tmp"
+  rm_rf "lib/tensorrt_rb/*.bundle"
+  rm_rf "lib/tensorrt_rb/*.so"
+  rm_f Dir.glob("*.gem")
+end

data/ext/tensorrt_rb/extconf.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "mkmf-rice"
+
+# TensorRT/CUDA paths for Linux (x86_64 and aarch64)
+case RUBY_PLATFORM
+when /x86_64-linux/
+  tensorrt_include = ENV["TENSORRT_INCLUDE"] || "/usr/include/x86_64-linux-gnu"
+  tensorrt_lib = ENV["TENSORRT_LIB"] || "/usr/lib/x86_64-linux-gnu"
+  cuda_include = ENV["CUDA_INCLUDE"] || "/usr/local/cuda/include"
+  cuda_lib = ENV["CUDA_LIB"] || "/usr/local/cuda/lib64"
+when /aarch64-linux/
+  tensorrt_include = ENV["TENSORRT_INCLUDE"] || "/usr/include/aarch64-linux-gnu"
+  tensorrt_lib = ENV["TENSORRT_LIB"] || "/usr/lib/aarch64-linux-gnu"
+  cuda_include = ENV["CUDA_INCLUDE"] || "/usr/local/cuda/include"
+  cuda_lib = ENV["CUDA_LIB"] || "/usr/local/cuda/lib64"
+else
+  abort "Unsupported platform: #{RUBY_PLATFORM}. Only Linux x86_64 and aarch64 are supported."
+end
+
+$INCFLAGS << " -I#{tensorrt_include} -I#{cuda_include}"
+$LDFLAGS << " -L#{tensorrt_lib} -L#{cuda_lib}"
+$LDFLAGS << " -lnvinfer -lcudart"
+
+# C++17 for Rice
+$CXXFLAGS << " -std=c++17"
+
+create_makefile("tensorrt_rb/tensorrt_rb")

data/ext/tensorrt_rb/tensorrt_rb.cpp

@@ -0,0 +1,175 @@
+#include <rice/rice.hpp>
+#include <rice/stl.hpp>
+#include <NvInfer.h>
+#include <cuda_runtime.h>
+#include <fstream>
+#include <vector>
+#include <memory>
+
+using namespace Rice;
+
+class TRTLogger : public nvinfer1::ILogger {
+public:
+    bool verbose = false;
+
+    void log(Severity severity, const char* msg) noexcept override {
+        if (verbose || severity <= Severity::kWARNING) {
+            const char* level = "";
+            switch (severity) {
+                case Severity::kINTERNAL_ERROR: level = "INTERNAL_ERROR"; break;
+                case Severity::kERROR: level = "ERROR"; break;
+                case Severity::kWARNING: level = "WARNING"; break;
+                case Severity::kINFO: level = "INFO"; break;
+                case Severity::kVERBOSE: level = "VERBOSE"; break;
+            }
+            if (verbose || severity <= Severity::kWARNING) {
+                fprintf(stderr, "[TensorRT %s] %s\n", level, msg);
+            }
+        }
+    }
+};
+
+class TRTEngine {
+private:
+    TRTLogger logger;
+    std::unique_ptr<nvinfer1::IRuntime> runtime;
+    std::unique_ptr<nvinfer1::ICudaEngine> engine;
+    std::unique_ptr<nvinfer1::IExecutionContext> context;
+    std::vector<void*> bindings;
+    cudaStream_t stream;
+
+public:
+    TRTEngine(const std::string& engine_path, bool verbose = false) {
+        logger.verbose = verbose;
+
+        std::ifstream file(engine_path, std::ios::binary);
+
+        if (!file) {
+            throw std::runtime_error("Failed to open engine file: " + engine_path);
+        }
+
+        file.seekg(0, std::ios::end);
+        size_t size = file.tellg();
+        file.seekg(0, std::ios::beg);
+
+        std::vector<char> buffer(size);
+        file.read(buffer.data(), size);
+
+        runtime.reset(nvinfer1::createInferRuntime(logger));
+
+        if (!runtime) {
+            throw std::runtime_error("Failed to create TensorRT runtime");
+        }
+
+        engine.reset(runtime->deserializeCudaEngine(buffer.data(), size));
+        if (!engine) {
+            throw std::runtime_error("Failed to deserialize engine");
+        }
+
+        context.reset(engine->createExecutionContext());
+        if (!context) {
+            throw std::runtime_error("Failed to create execution context");
+        }
+
+        cudaError_t err = cudaStreamCreate(&stream);
+        if (err != cudaSuccess) {
+            throw std::runtime_error("Failed to create CUDA stream");
+        }
+
+        int num_tensors = engine->getNbIOTensors();
+
+        bindings.resize(num_tensors, nullptr);
+    }
+
+    ~TRTEngine() {
+        if (stream) {
+            cudaStreamDestroy(stream);
+        }
+    }
+
+    int num_io_tensors() const {
+        return engine->getNbIOTensors();
+    }
+
+    std::string get_tensor_name(int index) const {
+        return engine->getIOTensorName(index);
+    }
+
+    bool is_input(const std::string& name) const {
+        return engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT;
+    }
+
+    std::vector<int64_t> get_tensor_shape(const std::string& name) const {
+        auto dims = engine->getTensorShape(name.c_str());
+        std::vector<int64_t> shape;
+        for (int i = 0; i < dims.nbDims; i++) {
+            shape.push_back(dims.d[i]);
+        }
+        return shape;
+    }
+
+    size_t get_tensor_bytes(const std::string& name) const {
+        auto dims = engine->getTensorShape(name.c_str());
+        size_t bytes = sizeof(float);
+        for (int i = 0; i < dims.nbDims; i++) {
+            bytes *= dims.d[i];
+        }
+        return bytes;
+    }
+
+    void set_tensor_address(const std::string& name, uint64_t ptr) {
+        void* addr = reinterpret_cast<void*>(ptr);
+        context->setTensorAddress(name.c_str(), addr);
+
+        for (int i = 0; i < engine->getNbIOTensors(); i++) {
+            if (name == engine->getIOTensorName(i)) {
+                bindings[i] = addr;
+                break;
+            }
+        }
+    }
+
+    bool execute() {
+        return context->executeV2(bindings.data());
+    }
+
+    bool enqueue() {
+        return context->enqueueV3(stream);
+    }
+
+    void memcpy_htod_async(uint64_t dst, const float* src, size_t count) {
+        cudaMemcpyAsync(reinterpret_cast<void*>(dst), src, count * sizeof(float),
+                        cudaMemcpyHostToDevice, stream);
+    }
+
+    void memcpy_dtoh_async(float* dst, uint64_t src, size_t count) {
+        cudaMemcpyAsync(dst, reinterpret_cast<void*>(src), count * sizeof(float),
+                        cudaMemcpyDeviceToHost, stream);
+    }
+
+    void stream_synchronize() {
+        cudaStreamSynchronize(stream);
+    }
+
+    uint64_t get_stream() const {
+        return reinterpret_cast<uint64_t>(stream);
+    }
+};
+
+extern "C" void Init_tensorrt_rb() {
+    Module rb_mTensorRT = define_module("TensorRT");
+
+    define_class_under<TRTEngine>(rb_mTensorRT, "Engine")
+        .define_constructor(Constructor<TRTEngine, const std::string&, bool>(),
+            Arg("engine_path"), Arg("verbose") = false)
+        .define_method("num_io_tensors", &TRTEngine::num_io_tensors)
+        .define_method("get_tensor_name", &TRTEngine::get_tensor_name)
+        .define_method("is_input?", &TRTEngine::is_input)
+        .define_method("get_tensor_shape", &TRTEngine::get_tensor_shape)
+        .define_method("get_tensor_bytes", &TRTEngine::get_tensor_bytes)
+        .define_method("set_tensor_address", &TRTEngine::set_tensor_address)
+        .define_method("execute", &TRTEngine::execute)
+        .define_method("enqueue", &TRTEngine::enqueue)
+        .define_method("stream_synchronize", &TRTEngine::stream_synchronize)
+        .define_method("get_stream", &TRTEngine::get_stream);
+}

data/lib/tensorrt/cuda.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require "ffi"
+
+module TensorRT
+  module CUDA
+    extend FFI::Library
+
+    CUDA_LIBS = %w[
+      libcudart.so
+      libcudart.so.12
+      libcudart.so.11
+    ].freeze
+
+    begin
+      ffi_lib CUDA_LIBS
+    rescue LoadError => e
+      raise LoadError, "Could not load CUDA runtime library. Tried: #{CUDA_LIBS.join(', ')}. Error: #{e.message}"
+    end
+
+    MEMCPY_HOST_TO_DEVICE = 1
+    MEMCPY_DEVICE_TO_HOST = 2
+
+    attach_function :cudaMalloc, [:pointer, :size_t], :int
+    attach_function :cudaFree, [:pointer], :int
+    attach_function :cudaMemcpy, [:pointer, :pointer, :size_t, :int], :int
+    attach_function :cudaMemcpyAsync, [:pointer, :pointer, :size_t, :int, :pointer], :int
+    attach_function :cudaDeviceSynchronize, [], :int
+    attach_function :cudaStreamSynchronize, [:pointer], :int
+
+    Error = Class.new(StandardError)
+
+    class << self
+      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
+
+      def free(ptr)
+        err = cudaFree(FFI::Pointer.new(ptr))
+
+        raise Error, "cudaFree failed with error #{err}" unless err.zero?
+      end
+
+      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
+
+      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
+
+      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
+
+      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
+
+      def synchronize
+        err = cudaDeviceSynchronize
+
+        raise Error, "cudaDeviceSynchronize failed with error #{err}" unless err.zero?
+      end
+
+      def stream_synchronize(stream)
+        err = cudaStreamSynchronize(FFI::Pointer.new(stream))
+
+        raise Error, "cudaStreamSynchronize failed with error #{err}" unless err.zero?
+      end
+    end
+
+    freeze
+  end
+end

data/lib/tensorrt.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require "tensorrt_rb/tensorrt_rb"
+
+module TensorRT
+  VERSION = "0.1.0"
+
+  autoload :CUDA, 'tensorrt/cuda'
+end

data/tensorrt.gemspec

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+Gem::Specification.new do |spec|
+  spec.name          = "tensorrt"
+  spec.version       = "0.1.0"
+  spec.author        = "Pete Mtasyburka"
+  spec.email         = ["pete@docuseal.com"]
+
+  spec.summary       = "Minimal TensorRT bindings for Ruby using Rice"
+  spec.description   = "Minimal Ruby bindings for NVIDIA TensorRT inference using Rice"
+  spec.homepage      = "https://github.com/docusealco/tensorrt-rb"
+  spec.license       = "Apache-2.0"
+  spec.required_ruby_version = ">= 3.0.0"
+
+  spec.files = Dir["lib/**/*.rb", "ext/**/*.{cpp,hpp,rb}", "*.gemspec", "Rakefile", "README.md"]
+  spec.require_paths = ["lib"]
+  spec.extensions    = ["ext/tensorrt_rb/extconf.rb"]
+
+  spec.add_dependency "rice", ">= 4.7"
+  spec.add_dependency "ffi", "~> 1.0"
+
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "rake-compiler", "~> 1.2"
+end

metadata

@@ -0,0 +1,104 @@
+--- !ruby/object:Gem::Specification
+name: tensorrt
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Pete Mtasyburka
+bindir: bin
+cert_chain: []
+date: 2025-12-28 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: rice
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '4.7'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '4.7'
+- !ruby/object:Gem::Dependency
+  name: ffi
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: rake-compiler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
+description: Minimal Ruby bindings for NVIDIA TensorRT inference using Rice
+email:
+- pete@docuseal.com
+executables: []
+extensions:
+- ext/tensorrt_rb/extconf.rb
+extra_rdoc_files: []
+files:
+- README.md
+- Rakefile
+- ext/tensorrt_rb/extconf.rb
+- ext/tensorrt_rb/tensorrt_rb.cpp
+- lib/tensorrt.rb
+- lib/tensorrt/cuda.rb
+- tensorrt.gemspec
+homepage: https://github.com/docusealco/tensorrt-rb
+licenses:
+- Apache-2.0
+metadata: {}
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.0.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.2
+specification_version: 4
+summary: Minimal TensorRT bindings for Ruby using Rice
+test_files: []