ignis 0.0.1

data/lib/nvruby/jit/kernel_module.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+module Ignis
+  module JIT
+    # Represents a loaded CUDA module on a specific device
+    # Manages the CUmodule handle and CUfunction extraction
+    # Handles automatic cleanup via finalizer
+    class KernelModule
+      # @return [CompiledKernel] The compiled kernel this module was loaded from
+      attr_reader :compiled_kernel
+
+      # @return [Integer] Device ID this module is loaded on
+      attr_reader :device_id
+
+      # @return [FFI::Pointer] The CUmodule handle
+      attr_reader :module_handle
+
+      # @return [FFI::Pointer] The CUfunction handle
+      attr_reader :function_handle
+
+      # @return [Time] When this module was loaded
+      attr_reader :loaded_at
+
+      # Create a new KernelModule by loading compiled code onto a device
+      # @param compiled_kernel [CompiledKernel] The compiled kernel to load
+      # @param device_id [Integer] Device to load onto
+      # @raise [CudaDriverError] If loading fails
+      def initialize(compiled_kernel, device_id: 0)
+        @compiled_kernel = compiled_kernel
+        @device_id = device_id
+        @loaded_at = Time.now
+        @destroyed = false
+        @mutex = Mutex.new
+
+        load_module!
+        extract_function!
+
+        setup_finalizer
+      end
+
+      # Get the kernel function name
+      # @return [String]
+      def kernel_name
+        @compiled_kernel.kernel_name
+      end
+
+      # Check if this module has been destroyed
+      # @return [Boolean]
+      def destroyed?
+        @mutex.synchronize { @destroyed }
+      end
+
+      # Create a Kernel instance for execution
+      # @return [Kernel] Executable kernel wrapper
+      def to_kernel
+        Kernel.new(self)
+      end
+
+      # Unload the module from GPU memory
+      # @return [void]
+      def destroy!
+        @mutex.synchronize do
+          return if @destroyed
+
+          if @module_handle && !@module_handle.null?
+            DriverAPIBindings.unload_module(@module_handle)
+          end
+
+          @module_handle = nil
+          @function_handle = nil
+          @destroyed = true
+
+          Ignis.logger.debug("Unloaded kernel module: #{kernel_name} from device #{@device_id}")
+        end
+      end
+
+      # Get a string representation
+      # @return [String]
+      def to_s
+        "#<Ignis::JIT::KernelModule #{kernel_name} device=#{@device_id}>"
+      end
+
+      # Get detailed inspection
+      # @return [String]
+      def inspect
+        status = @destroyed ? "destroyed" : "loaded"
+        "#<Ignis::JIT::KernelModule:0x#{object_id.to_s(16)} " \
+          "kernel=#{kernel_name.inspect} " \
+          "device=#{@device_id} " \
+          "status=#{status}>"
+      end
+
+      private
+
+      # Load the CUBIN data as a module
+      # @return [void]
+      def load_module!
+        Ignis.set_device(@device_id)
+        @module_handle = DriverAPIBindings.load_module_data(@compiled_kernel.cubin_data)
+        Ignis.logger.debug("Loaded module from CUBIN (#{@compiled_kernel.cubin_size} bytes)")
+      end
+
+      # Extract the kernel function from the module
+      # @return [void]
+      def extract_function!
+        @function_handle = DriverAPIBindings.get_module_function(@module_handle, kernel_name)
+        Ignis.logger.debug("Extracted function: #{kernel_name}")
+      end
+
+      # Setup finalizer for automatic cleanup
+      # @return [void]
+      def setup_finalizer
+        destructor = destructor_proc(@module_handle)
+        ObjectSpace.define_finalizer(self, destructor)
+      end
+
+      # Create a destructor proc that doesn't capture self
+      # @param module_handle [FFI::Pointer] Module handle to unload
+      # @return [Proc]
+      def destructor_proc(module_handle)
+        proc do
+          if module_handle && !module_handle.null?
+            begin
+              DriverAPIBindings.unload_module(module_handle)
+            rescue StandardError
+              nil
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/jit/kernels/activations.rb

@@ -0,0 +1,179 @@
+# frozen_string_literal: true
+
+module Ignis
+  module JIT
+    module Kernels
+      # Activation function CUDA kernels for AI training and inference.
+      # Each kernel has a forward and backward variant for autograd compatibility.
+      # All kernels are compiled via NVRTC on first use and cached.
+      module Activations
+        class << self
+          # ReLU forward: max(0, x)
+          # @param n [Integer] number of elements
+          # @return [Ignis::JIT::Kernel]
+          def relu_forward(n)
+            source = <<~CUDA
+              extern "C" __global__
+              void relu_forward(const float* __restrict__ input,
+                                float* __restrict__ output,
+                                const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  output[idx] = fmaxf(input[idx], 0.0f);
+                }
+              }
+            CUDA
+            compile_cached(source, "relu_forward")
+          end
+
+          # ReLU backward: grad * (x > 0 ? 1 : 0)
+          # @return [Ignis::JIT::Kernel]
+          def relu_backward
+            source = <<~CUDA
+              extern "C" __global__
+              void relu_backward(const float* __restrict__ grad_output,
+                                 const float* __restrict__ input,
+                                 float* __restrict__ grad_input,
+                                 const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  grad_input[idx] = input[idx] > 0.0f ? grad_output[idx] : 0.0f;
+                }
+              }
+            CUDA
+            compile_cached(source, "relu_backward")
+          end
+
+          # GELU forward: x * Φ(x) using tanh approximation
+          # GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
+          # @return [Ignis::JIT::Kernel]
+          def gelu_forward
+            source = <<~CUDA
+              extern "C" __global__
+              void gelu_forward(const float* __restrict__ input,
+                                float* __restrict__ output,
+                                const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  float x = input[idx];
+                  float cdf = 0.5f * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x)));
+                  output[idx] = x * cdf;
+                }
+              }
+            CUDA
+            compile_cached(source, "gelu_forward")
+          end
+
+          # GELU backward: Φ(x) + x * φ(x)
+          # @return [Ignis::JIT::Kernel]
+          def gelu_backward
+            source = <<~CUDA
+              extern "C" __global__
+              void gelu_backward(const float* __restrict__ grad_output,
+                                 const float* __restrict__ input,
+                                 float* __restrict__ grad_input,
+                                 const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  float x = input[idx];
+                  float s = 0.7978845608f * (x + 0.044715f * x * x * x);
+                  float tanh_s = tanhf(s);
+                  float cdf = 0.5f * (1.0f + tanh_s);
+                  float pdf_term = 0.5f * (1.0f - tanh_s * tanh_s) * 0.7978845608f * (1.0f + 3.0f * 0.044715f * x * x);
+                  grad_input[idx] = grad_output[idx] * (cdf + x * pdf_term);
+                }
+              }
+            CUDA
+            compile_cached(source, "gelu_backward")
+          end
+
+          # SiLU forward: x * sigmoid(x)
+          # @return [Ignis::JIT::Kernel]
+          def silu_forward
+            source = <<~CUDA
+              extern "C" __global__
+              void silu_forward(const float* __restrict__ input,
+                                float* __restrict__ output,
+                                const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  float x = input[idx];
+                  float sig = 1.0f / (1.0f + expf(-x));
+                  output[idx] = x * sig;
+                }
+              }
+            CUDA
+            compile_cached(source, "silu_forward")
+          end
+
+          # SiLU backward: sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x))
+          # = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
+          # @return [Ignis::JIT::Kernel]
+          def silu_backward
+            source = <<~CUDA
+              extern "C" __global__
+              void silu_backward(const float* __restrict__ grad_output,
+                                 const float* __restrict__ input,
+                                 float* __restrict__ grad_input,
+                                 const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  float x = input[idx];
+                  float sig = 1.0f / (1.0f + expf(-x));
+                  grad_input[idx] = grad_output[idx] * (sig * (1.0f + x * (1.0f - sig)));
+                }
+              }
+            CUDA
+            compile_cached(source, "silu_backward")
+          end
+
+          # Sigmoid forward: 1 / (1 + exp(-x))
+          # @return [Ignis::JIT::Kernel]
+          def sigmoid_forward
+            source = <<~CUDA
+              extern "C" __global__
+              void sigmoid_forward(const float* __restrict__ input,
+                                   float* __restrict__ output,
+                                   const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  output[idx] = 1.0f / (1.0f + expf(-input[idx]));
+                }
+              }
+            CUDA
+            compile_cached(source, "sigmoid_forward")
+          end
+
+          # Sigmoid backward: sigmoid(x) * (1 - sigmoid(x)) * grad
+          # @return [Ignis::JIT::Kernel]
+          def sigmoid_backward
+            source = <<~CUDA
+              extern "C" __global__
+              void sigmoid_backward(const float* __restrict__ grad_output,
+                                    const float* __restrict__ output,
+                                    float* __restrict__ grad_input,
+                                    const int n) {
+                int idx = blockIdx.x * blockDim.x + threadIdx.x;
+                if (idx < n) {
+                  float s = output[idx];
+                  grad_input[idx] = grad_output[idx] * s * (1.0f - s);
+                }
+              }
+            CUDA
+            compile_cached(source, "sigmoid_backward")
+          end
+
+          private
+
+          # @param source [String] CUDA source code
+          # @param name [String] kernel function name
+          # @param device_id [Integer] target device
+          # @return [Ignis::JIT::Kernel]
+          def compile_cached(source, name, device_id: 0)
+            Ignis::JIT::Compiler.compile(source, name, device_id: device_id)
+          end
+        end
+      end
+    end
+  end
+end