ignis 0.0.1

data/lib/nvruby/epilogues.rb

@@ -0,0 +1,438 @@
+# frozen_string_literal: true
+
+module Ignis
+  # Advanced fused epilogues for GPU operations
+  # Provides GELU, ReLU, SiLU, Bias addition as fused CUDA kernels
+  #
+  # @example Apply GELU activation
+  #   output = Ignis::Epilogues.gelu(input)
+  #
+  # @example Fused GEMM + GELU + Bias
+  #   output = Ignis::Epilogues.gemm_gelu_bias(a, b, bias)
+  #
+  module Epilogues
+    # GELU approximation constant
+    GELU_COEF_A = 0.7978845608028654
+    GELU_COEF_B = 0.044715
+
+    # JIT CUDA kernels for epilogues
+    module Kernels
+      GELU_KERNEL = <<~CUDA
+        extern "C" __global__ void gelu_forward(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                float x = input[idx];
+                // Approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+                float x3 = x * x * x;
+                float tanh_arg = 0.7978845608f * (x + 0.044715f * x3);
+                output[idx] = 0.5f * x * (1.0f + tanhf(tanh_arg));
+            }
+        }
+      CUDA
+
+      GELU_EXACT_KERNEL = <<~CUDA
+        extern "C" __global__ void gelu_exact_forward(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                float x = input[idx];
+                // Exact: x * 0.5 * (1 + erf(x / sqrt(2)))
+                output[idx] = x * 0.5f * (1.0f + erff(x * 0.7071067811865476f));
+            }
+        }
+      CUDA
+
+      SILU_KERNEL = <<~CUDA
+        extern "C" __global__ void silu_forward(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                float x = input[idx];
+                // SiLU: x * sigmoid(x) = x / (1 + exp(-x))
+                output[idx] = x / (1.0f + expf(-x));
+            }
+        }
+      CUDA
+
+      RELU_KERNEL = <<~CUDA
+        extern "C" __global__ void relu_forward(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                output[idx] = fmaxf(0.0f, input[idx]);
+            }
+        }
+      CUDA
+
+      LEAKY_RELU_KERNEL = <<~CUDA
+        extern "C" __global__ void leaky_relu_forward(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            int n,
+            float negative_slope
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                float x = input[idx];
+                output[idx] = x > 0.0f ? x : x * negative_slope;
+            }
+        }
+      CUDA
+
+      BIAS_ADD_KERNEL = <<~CUDA
+        extern "C" __global__ void bias_add(
+            const float* __restrict__ input,
+            const float* __restrict__ bias,
+            float* __restrict__ output,
+            int rows,
+            int cols
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            int total = rows * cols;
+            if (idx < total) {
+                int col = idx % cols;
+                output[idx] = input[idx] + bias[col];
+            }
+        }
+      CUDA
+
+      GELU_BIAS_KERNEL = <<~CUDA
+        extern "C" __global__ void gelu_bias_forward(
+            const float* __restrict__ input,
+            const float* __restrict__ bias,
+            float* __restrict__ output,
+            int rows,
+            int cols
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            int total = rows * cols;
+            if (idx < total) {
+                int col = idx % cols;
+                float x = input[idx] + bias[col];
+                float x3 = x * x * x;
+                float tanh_arg = 0.7978845608f * (x + 0.044715f * x3);
+                output[idx] = 0.5f * x * (1.0f + tanhf(tanh_arg));
+            }
+        }
+      CUDA
+
+      SILU_BIAS_KERNEL = <<~CUDA
+        extern "C" __global__ void silu_bias_forward(
+            const float* __restrict__ input,
+            const float* __restrict__ bias,
+            float* __restrict__ output,
+            int rows,
+            int cols
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            int total = rows * cols;
+            if (idx < total) {
+                int col = idx % cols;
+                float x = input[idx] + bias[col];
+                output[idx] = x / (1.0f + expf(-x));
+            }
+        }
+      CUDA
+
+      RESIDUAL_ADD_KERNEL = <<~CUDA
+        extern "C" __global__ void residual_add(
+            const float* __restrict__ input,
+            const float* __restrict__ residual,
+            float* __restrict__ output,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                output[idx] = input[idx] + residual[idx];
+            }
+        }
+      CUDA
+
+      SCALE_KERNEL = <<~CUDA
+        extern "C" __global__ void scale(
+            const float* __restrict__ input,
+            float* __restrict__ output,
+            float scale_factor,
+            int n
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx < n) {
+                output[idx] = input[idx] * scale_factor;
+            }
+        }
+      CUDA
+    end
+
+    class << self
+      # Apply GELU activation (approximation)
+      #
+      # @param input [NvArray] Input tensor
+      # @param out [NvArray, nil] Output tensor (optional)
+      # @return [NvArray] Output with GELU applied
+      def gelu(input, out: nil)
+        apply_unary(input, out, :gelu, Kernels::GELU_KERNEL, "gelu_forward")
+      end
+
+      # Apply exact GELU activation
+      #
+      # @param input [NvArray] Input tensor
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def gelu_exact(input, out: nil)
+        apply_unary(input, out, :gelu_exact, Kernels::GELU_EXACT_KERNEL, "gelu_exact_forward")
+      end
+
+      # Apply SiLU (Swish) activation
+      #
+      # @param input [NvArray] Input tensor
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def silu(input, out: nil)
+        apply_unary(input, out, :silu, Kernels::SILU_KERNEL, "silu_forward")
+      end
+
+      # Apply ReLU activation
+      #
+      # @param input [NvArray] Input tensor
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def relu(input, out: nil)
+        apply_unary(input, out, :relu, Kernels::RELU_KERNEL, "relu_forward")
+      end
+
+      # Apply Leaky ReLU activation
+      #
+      # @param input [NvArray] Input tensor
+      # @param negative_slope [Float] Slope for negative values
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def leaky_relu(input, negative_slope: 0.01, out: nil)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        n = input.size
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(:leaky_relu, Kernels::LEAKY_RELU_KERNEL, "leaky_relu_forward")
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, out.device_ffi_ptr, n, negative_slope]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      # Add bias to tensor
+      #
+      # @param input [NvArray] Input tensor (rows x cols)
+      # @param bias [NvArray] Bias vector (cols)
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def bias_add(input, bias, out: nil)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        shape = input.shape
+        rows = shape.size == 1 ? 1 : shape[0]
+        cols = shape.size == 1 ? shape[0] : shape[1]
+        n = rows * cols
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(:bias_add, Kernels::BIAS_ADD_KERNEL, "bias_add")
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, bias.device_ffi_ptr, out.device_ffi_ptr, rows, cols]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      # Fused GELU + Bias
+      #
+      # @param input [NvArray] Input tensor
+      # @param bias [NvArray] Bias vector
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def gelu_bias(input, bias, out: nil)
+        apply_fused_bias(input, bias, out, :gelu_bias, Kernels::GELU_BIAS_KERNEL, "gelu_bias_forward")
+      end
+
+      # Fused SiLU + Bias
+      #
+      # @param input [NvArray] Input tensor
+      # @param bias [NvArray] Bias vector
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def silu_bias(input, bias, out: nil)
+        apply_fused_bias(input, bias, out, :silu_bias, Kernels::SILU_BIAS_KERNEL, "silu_bias_forward")
+      end
+
+      # Residual addition
+      #
+      # @param input [NvArray] Input tensor
+      # @param residual [NvArray] Residual tensor
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def residual_add(input, residual, out: nil)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        n = input.size
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(:residual_add, Kernels::RESIDUAL_ADD_KERNEL, "residual_add")
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, residual.device_ffi_ptr, out.device_ffi_ptr, n]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      # Scale tensor by factor
+      #
+      # @param input [NvArray] Input tensor
+      # @param factor [Float] Scale factor
+      # @param out [NvArray, nil] Output tensor
+      # @return [NvArray]
+      def scale(input, factor, out: nil)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        n = input.size
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(:scale, Kernels::SCALE_KERNEL, "scale")
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, out.device_ffi_ptr, factor, n]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      # Fused GEMM + epilogue
+      #
+      # @param a [NvArray] First matrix
+      # @param b [NvArray] Second matrix
+      # @param epilogue [Symbol] :gelu, :relu, :silu
+      # @param bias [NvArray, nil] Optional bias
+      # @return [NvArray]
+      def gemm_epilogue(a, b, epilogue:, bias: nil)
+        # Perform GEMM
+        c = Ignis::LinAlg.matmul(a, b)
+
+        # Apply epilogue
+        result = case epilogue
+                 when :gelu
+                   bias ? gelu_bias(c, bias) : gelu(c)
+                 when :relu
+                   temp = bias ? bias_add(c, bias) : c
+                   relu(temp)
+                 when :silu
+                   bias ? silu_bias(c, bias) : silu(c)
+                 else
+                   bias ? bias_add(c, bias) : c
+                 end
+
+        result
+      end
+
+      private
+
+      def apply_unary(input, out, name, kernel_code, kernel_name)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        n = input.size
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(name, kernel_code, kernel_name)
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, out.device_ffi_ptr, n]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      def apply_fused_bias(input, bias, out, name, kernel_code, kernel_name)
+        CUDA::RuntimeAPI.ensure_loaded!
+
+        shape = input.shape
+        rows = shape.size == 1 ? 1 : shape[0]
+        cols = shape.size == 1 ? shape[0] : shape[1]
+        n = rows * cols
+        device = input.respond_to?(:device_index) ? input.device_index : 0
+
+        out ||= Ignis::NvArray.zeros(input.shape, dtype: input.dtype, device: device)
+
+        kernel = get_kernel(name, kernel_code, kernel_name)
+
+        block_size = 256
+        grid_size = (n + block_size - 1) / block_size
+
+        kernel.launch(
+          grid: [grid_size, 1, 1],
+          block: [block_size, 1, 1],
+          args: [input.device_ffi_ptr, bias.device_ffi_ptr, out.device_ffi_ptr, rows, cols]
+        )
+
+        CUDA::RuntimeAPI.cudaDeviceSynchronize
+        out
+      end
+
+      def get_kernel(name, kernel_code, kernel_name)
+        @kernels ||= {}
+        @kernels[name] ||= begin
+          compiler = CUDA::JITCompiler.new
+          compiler.compile(kernel_code, kernel_name)
+        end
+      end
+    end
+  end
+end