ignis-numerics 0.0.1

data/lib/nvruby/mathdx/fft_kernel.rb

@@ -0,0 +1,258 @@
+# frozen_string_literal: true
+
+module Ignis
+  module MathDx
+    # Device-side FFT kernel using cuFFTDx patterns
+    # Generates and compiles CUDA C++ code for thread block FFT operations
+    #
+    # cuFFTDx enables embedding FFT operations inside CUDA kernels, allowing
+    # fusion with other operations to reduce memory bandwidth.
+    #
+    # @example Basic FFT kernel
+    #   kernel = FftKernel.new(size: 64, dtype: :complex64, direction: :forward)
+    #   kernel.compile!
+    #   kernel.execute(input, output)
+    class FftKernel
+      # @return [Integer] FFT size
+      attr_reader :size
+
+      # @return [Symbol] Data type
+      attr_reader :dtype
+
+      # @return [Symbol] FFT direction
+      attr_reader :direction
+
+      # @return [Integer] Elements per thread
+      attr_reader :elements_per_thread
+
+      # @return [Boolean] Whether kernel is compiled
+      attr_reader :compiled
+
+      # Supported FFT sizes (powers of 2)
+      SUPPORTED_SIZES = [16, 32, 64, 128, 256, 512, 1024].freeze
+
+      # Initialize FFT kernel configuration
+      # @param size [Integer] FFT size (power of 2, <= 1024)
+      # @param dtype [Symbol] Data type (:complex64, :complex128)
+      # @param direction [Symbol] Direction (:forward, :inverse)
+      # @param elements_per_thread [Integer] Elements per thread (default: 8)
+      def initialize(size:, dtype: :complex64, direction: :forward, elements_per_thread: 8)
+        validate_size!(size)
+        validate_dtype!(dtype)
+        validate_direction!(direction)
+
+        @size = size
+        @dtype = dtype
+        @direction = direction
+        @elements_per_thread = elements_per_thread
+        @compiled = false
+        @kernel = nil
+      end
+
+      # Compile the FFT kernel
+      # @param device_id [Integer] Target GPU device
+      # @return [self]
+      def compile!(device_id: 0)
+        source = generate_source
+        @kernel = Ignis::JIT::Compiler.compile(
+          source,
+          "cufftdx_fft",
+          device_id: device_id,
+          options: nvrtc_options
+        )
+        @compiled = true
+        self
+      end
+
+      # Execute the FFT kernel
+      # @param input [NvArray] Input array (complex)
+      # @param output [NvArray, nil] Output array (created if nil)
+      # @param batch [Integer] Number of batched FFTs
+      # @param stream [CUDA::Stream, nil] CUDA stream
+      # @return [NvArray] Result array
+      def execute(input, output: nil, batch: 1, stream: nil)
+        raise StateError, "Kernel not compiled. Call compile! first." unless @compiled
+
+        validate_execution_input!(input)
+
+        # Ensure input is on device
+        input_dev = input.on_device? ? input : input.to_device
+
+        # Create output if needed
+        output_dev = if output
+                       output.on_device? ? output : output.to_device
+                     else
+                       NvArray.zeros(input.shape, dtype: @dtype, device: input_dev.device_index).to_device
+                     end
+
+        # Calculate grid dimensions
+        threads_per_fft = @size / @elements_per_thread
+        blocks = batch
+
+        # Launch kernel
+        @kernel.launch(
+          grid: [blocks],
+          block: [threads_per_fft],
+          shared_memory: shared_memory_size,
+          args: [
+            input_dev.device_ptr,
+            output_dev.device_ptr,
+            @size,
+            batch
+          ],
+          stream: stream
+        )
+
+        output_dev
+      end
+
+      # Release kernel resources
+      # @return [void]
+      def destroy!
+        @kernel = nil
+        @compiled = false
+      end
+
+      private
+
+      # Validate FFT size
+      def validate_size!(size)
+        unless size.positive? && (size & (size - 1)).zero? && size <= 1024
+          raise ArgumentError, "FFT size must be power of 2 <= 1024, got #{size}"
+        end
+      end
+
+      # Validate data type
+      def validate_dtype!(dtype)
+        unless %i[complex64 complex128].include?(dtype)
+          raise UnsupportedDTypeError.new(dtype, operation: "cuFFTDx FFT")
+        end
+      end
+
+      # Validate direction
+      def validate_direction!(direction)
+        unless %i[forward inverse].include?(direction)
+          raise ArgumentError, "Direction must be :forward or :inverse, got #{direction}"
+        end
+      end
+
+      # Validate execution input
+      def validate_execution_input!(input)
+        unless input.is_a?(NvArray)
+          raise ArgumentError, "Expected NvArray input"
+        end
+
+        unless input.dtype == @dtype
+          raise ArgumentError, "Input dtype must be #{@dtype}, got #{input.dtype}"
+        end
+      end
+
+      # Calculate shared memory size
+      # @return [Integer] Bytes of shared memory
+      def shared_memory_size
+        bytes_per_element = @dtype == :complex64 ? 8 : 16
+        @size * bytes_per_element
+      end
+
+      # Get NVRTC compilation options
+      # @return [Array<String>]
+      def nvrtc_options
+        opts = [
+          "-I#{MathDx.include_path}",
+          "--std=c++17",
+          "-DNDEBUG"
+        ]
+
+        cuda_include = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.0/include"
+        opts << "-I#{cuda_include}" if File.exist?(cuda_include)
+
+        opts
+      end
+
+      # Generate CUDA C++ source code
+      # @return [String]
+      def generate_source
+        complex_type = @dtype == :complex64 ? "float2" : "double2"
+        base_type = @dtype == :complex64 ? "float" : "double"
+        sign = @direction == :forward ? -1 : 1
+
+        <<~CUDA
+          #include <cuda_runtime.h>
+          #include <math.h>
+          
+          // Complex multiplication
+          __device__ inline #{complex_type} cmul(#{complex_type} a, #{complex_type} b) {
+              return make_#{complex_type}(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
+          }
+          
+          // Complex addition
+          __device__ inline #{complex_type} cadd(#{complex_type} a, #{complex_type} b) {
+              return make_#{complex_type}(a.x + b.x, a.y + b.y);
+          }
+          
+          // Complex subtraction
+          __device__ inline #{complex_type} csub(#{complex_type} a, #{complex_type} b) {
+              return make_#{complex_type}(a.x - b.x, a.y - b.y);
+          }
+          
+          // Twiddle factor
+          __device__ inline #{complex_type} twiddle(int k, int N) {
+              #{base_type} angle = #{sign} * 2.0#{base_type == "double" ? "" : "f"} * M_PI * k / N;
+              return make_#{complex_type}(cos#{base_type == "double" ? "" : "f"}(angle), sin#{base_type == "double" ? "" : "f"}(angle));
+          }
+          
+          // Cooley-Tukey radix-2 FFT in shared memory
+          extern "C" __global__ void cufftdx_fft(
+              const #{complex_type}* __restrict__ input,
+              #{complex_type}* __restrict__ output,
+              int N,
+              int batch
+          ) {
+              extern __shared__ #{complex_type} shared[];
+              
+              int batch_idx = blockIdx.x;
+              int tid = threadIdx.x;
+              int threads = blockDim.x;
+              
+              // Load data into shared memory with bit-reversal
+              for (int i = tid; i < N; i += threads) {
+                  // Bit-reversal permutation
+                  int rev = 0;
+                  int val = i;
+                  for (int j = 1; j < N; j <<= 1) {
+                      rev = (rev << 1) | (val & 1);
+                      val >>= 1;
+                  }
+                  shared[i] = input[batch_idx * N + rev];
+              }
+              __syncthreads();
+              
+              // Iterative Cooley-Tukey FFT
+              for (int s = 1; s < N; s <<= 1) {
+                  int m = s << 1;
+                  for (int k = tid; k < N / 2; k += threads) {
+                      int j = k / s;
+                      int i = k % s;
+                      int idx1 = j * m + i;
+                      int idx2 = idx1 + s;
+                      
+                      #{complex_type} w = twiddle(i, m);
+                      #{complex_type} t = cmul(w, shared[idx2]);
+                      #{complex_type} u = shared[idx1];
+                      
+                      shared[idx1] = cadd(u, t);
+                      shared[idx2] = csub(u, t);
+                  }
+                  __syncthreads();
+              }
+              
+              // Write output
+              for (int i = tid; i < N; i += threads) {
+                  #{@direction == :inverse ? "output[batch_idx * N + i] = make_#{complex_type}(shared[i].x / N, shared[i].y / N);" : "output[batch_idx * N + i] = shared[i];"}
+              }
+          }
+        CUDA
+      end
+    end
+  end
+end

data/lib/nvruby/mathdx/gemm_kernel.rb

@@ -0,0 +1,293 @@
+# frozen_string_literal: true
+
+module Ignis
+  module MathDx
+    # Device-side GEMM kernel using cuBLASDx
+    # Generates and compiles CUDA C++ code that uses cuBLASDx for thread block GEMM
+    #
+    # cuBLASDx enables embedding GEMM operations inside CUDA kernels, allowing
+    # fusion with other operations (epilogs) to reduce memory bandwidth.
+    #
+    # @example Basic GEMM kernel
+    #   kernel = GemmKernel.new(m: 64, n: 64, k: 64, dtype: :float16)
+    #   kernel.compile!
+    #   kernel.execute(a, b, c)
+    class GemmKernel
+      # @return [Integer] Matrix M dimension
+      attr_reader :m
+
+      # @return [Integer] Matrix N dimension
+      attr_reader :n
+
+      # @return [Integer] Matrix K dimension
+      attr_reader :k
+
+      # @return [Symbol] Data type
+      attr_reader :dtype
+
+      # @return [Symbol, nil] Epilog operation
+      attr_reader :epilog
+
+      # @return [Integer] Thread block size
+      attr_reader :block_size
+
+      # @return [Boolean] Whether kernel is compiled
+      attr_reader :compiled
+
+      # Initialize GEMM kernel configuration
+      # @param m [Integer] Output rows (must be power of 2, <= 128)
+      # @param n [Integer] Output columns (must be power of 2, <= 128)
+      # @param k [Integer] Inner dimension (must be power of 2)
+      # @param dtype [Symbol] Data type (:float16, :float32, :float64)
+      # @param epilog [Symbol, nil] Epilog (:relu, :gelu, :sigmoid, :tanh, nil)
+      # @param block_size [Integer] Threads per block
+      def initialize(m:, n:, k:, dtype: :float32, epilog: nil, block_size: 128)
+        validate_dimensions!(m, n, k)
+        validate_dtype!(dtype)
+        validate_epilog!(epilog) if epilog
+
+        @m = m
+        @n = n
+        @k = k
+        @dtype = dtype
+        @epilog = epilog
+        @block_size = block_size
+        @compiled = false
+        @kernel = nil
+      end
+
+      # Compile the GEMM kernel
+      # @param device_id [Integer] Target GPU device
+      # @return [self]
+      def compile!(device_id: 0)
+        source = generate_source
+        @kernel = Ignis::JIT::Compiler.compile(
+          source,
+          "cublasdx_gemm",
+          device_id: device_id,
+          options: nvrtc_options
+        )
+        @compiled = true
+        self
+      end
+
+      # Execute the GEMM kernel: C = alpha * A @ B + beta * C
+      # @param a [NvArray] Input matrix A (M x K)
+      # @param b [NvArray] Input matrix B (K x N)
+      # @param c [NvArray] Output matrix C (M x N)
+      # @param alpha [Float] Scaling factor for A @ B
+      # @param beta [Float] Scaling factor for C
+      # @param stream [CUDA::Stream, nil] CUDA stream
+      # @return [NvArray] Result matrix C
+      def execute(a, b, c, alpha: 1.0, beta: 0.0, stream: nil)
+        raise StateError, "Kernel not compiled. Call compile! first." unless @compiled
+
+        validate_execution_inputs!(a, b, c)
+
+        # Ensure arrays are on device
+        a_dev = a.on_device? ? a : a.to_device
+        b_dev = b.on_device? ? b : b.to_device
+        c_dev = c.on_device? ? c : c.to_device
+
+        # Calculate grid dimensions
+        grid_m = (a_dev.shape[0] + @m - 1) / @m
+        grid_n = (b_dev.shape[1] + @n - 1) / @n
+
+        # Launch kernel
+        @kernel.launch(
+          grid: [grid_m, grid_n],
+          block: [@block_size],
+          args: [
+            a_dev.device_ptr,
+            b_dev.device_ptr,
+            c_dev.device_ptr,
+            a_dev.shape[0],         # M
+            b_dev.shape[1],         # N
+            a_dev.shape[1],         # K
+            a_dev.shape[1],         # lda
+            b_dev.shape[1],         # ldb
+            c_dev.shape[1],         # ldc
+            alpha,
+            beta
+          ],
+          stream: stream
+        )
+
+        c_dev
+      end
+
+      # Release kernel resources
+      # @return [void]
+      def destroy!
+        @kernel = nil
+        @compiled = false
+      end
+
+      private
+
+      # Validate dimensions
+      def validate_dimensions!(m, n, k)
+        [m, n, k].each do |dim|
+          unless dim.positive? && (dim & (dim - 1)).zero? && dim <= 128
+            raise ArgumentError, "Dimensions must be positive powers of 2 <= 128, got #{dim}"
+          end
+        end
+      end
+
+      # Validate data type
+      def validate_dtype!(dtype)
+        unless %i[float16 float32 float64].include?(dtype)
+          raise UnsupportedDTypeError.new(dtype, operation: "cuBLASDx GEMM")
+        end
+      end
+
+      # Validate epilog
+      def validate_epilog!(epilog)
+        unless %i[relu gelu sigmoid tanh].include?(epilog)
+          raise ArgumentError, "Unknown epilog: #{epilog}. Supported: relu, gelu, sigmoid, tanh"
+        end
+      end
+
+      # Validate execution inputs
+      def validate_execution_inputs!(a, b, c)
+        unless a.is_a?(NvArray) && b.is_a?(NvArray) && c.is_a?(NvArray)
+          raise ArgumentError, "Expected NvArray inputs"
+        end
+
+        unless a.dtype == @dtype && b.dtype == @dtype && c.dtype == @dtype
+          raise ArgumentError, "Array dtypes must match kernel dtype #{@dtype}"
+        end
+      end
+
+      # Get NVRTC compilation options
+      # @return [Array<String>]
+      def nvrtc_options
+        opts = [
+          "-I#{MathDx.include_path}",
+          "--std=c++17",
+          "-DNDEBUG"
+        ]
+
+        # Add CUDA include path
+        cuda_include = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.0/include"
+        opts << "-I#{cuda_include}" if File.exist?(cuda_include)
+
+        opts
+      end
+
+      # Generate CUDA C++ source code
+      # @return [String]
+      def generate_source
+        cuda_type = case @dtype
+                    when :float16 then "__half"
+                    when :float32 then "float"
+                    when :float64 then "double"
+                    end
+
+        epilog_code = generate_epilog_code
+
+        <<~CUDA
+          #include <cuda_fp16.h>
+          
+          // cuBLASDx fallback implementation when SDK not available
+          // Uses shared memory tiled GEMM algorithm
+          
+          #{epilog_code[:device_function]}
+          
+          extern "C" __global__ void cublasdx_gemm(
+              const #{cuda_type}* __restrict__ A,
+              const #{cuda_type}* __restrict__ B,
+              #{cuda_type}* __restrict__ C,
+              int M, int N, int K,
+              int lda, int ldb, int ldc,
+              #{cuda_type} alpha, #{cuda_type} beta
+          ) {
+              // Thread block GEMM with shared memory tiling
+              const int TILE_M = #{@m};
+              const int TILE_N = #{@n};
+              const int TILE_K = #{@k};
+              
+              __shared__ #{cuda_type} As[TILE_M][TILE_K];
+              __shared__ #{cuda_type} Bs[TILE_K][TILE_N];
+              
+              int bx = blockIdx.x;
+              int by = blockIdx.y;
+              int tx = threadIdx.x % TILE_N;
+              int ty = threadIdx.x / TILE_N;
+              
+              int row = bx * TILE_M + ty;
+              int col = by * TILE_N + tx;
+              
+              #{cuda_type} acc = 0;
+              
+              for (int t = 0; t < (K + TILE_K - 1) / TILE_K; ++t) {
+                  // Load tile of A into shared memory
+                  if (row < M && t * TILE_K + tx < K) {
+                      As[ty][tx] = A[row * lda + t * TILE_K + tx];
+                  } else {
+                      As[ty][tx] = 0;
+                  }
+                  
+                  // Load tile of B into shared memory
+                  if (t * TILE_K + ty < K && col < N) {
+                      Bs[ty][tx] = B[(t * TILE_K + ty) * ldb + col];
+                  } else {
+                      Bs[ty][tx] = 0;
+                  }
+                  
+                  __syncthreads();
+                  
+                  // Compute partial dot product
+                  #pragma unroll
+                  for (int k = 0; k < TILE_K; ++k) {
+                      acc += As[ty][k] * Bs[k][tx];
+                  }
+                  
+                  __syncthreads();
+              }
+              
+              // Write result with epilog
+              if (row < M && col < N) {
+                  #{cuda_type} result = alpha * acc + beta * C[row * ldc + col];
+                  #{epilog_code[:apply]}
+                  C[row * ldc + col] = result;
+              }
+          }
+        CUDA
+      end
+
+      # Generate epilog code
+      # @return [Hash] :device_function and :apply code
+      def generate_epilog_code
+        case @epilog
+        when :relu
+          {
+            device_function: "__device__ inline float relu(float x) { return fmaxf(x, 0.0f); }",
+            apply: "result = relu(result);"
+          }
+        when :gelu
+          {
+            device_function: "__device__ inline float gelu(float x) { " \
+                             "return 0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x))); }",
+            apply: "result = gelu(result);"
+          }
+        when :sigmoid
+          {
+            device_function: "__device__ inline float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); }",
+            apply: "result = sigmoid(result);"
+          }
+        when :tanh
+          {
+            device_function: "",
+            apply: "result = tanhf(result);"
+          }
+        else
+          { device_function: "", apply: "" }
+        end
+      end
+    end
+
+    # State error for kernel operations
+    class StateError < StandardError; end
+  end
+end

data/lib/nvruby/mathdx.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+# Ignis MathDx Module
+# Device-side BLAS/FFT extensions using cuBLASDx and cuFFTDx
+
+require_relative "mathdx/gemm_kernel"
+require_relative "mathdx/fft_kernel"
+
+module Ignis
+  # MathDx module for device-side extensions
+  # Provides high-performance fused kernel generation using cuBLASDx and cuFFTDx
+  #
+  # Note: MathDx libraries (cuBLASDx, cuFFTDx) are header-only C++ template libraries
+  # that generate optimized device code at compile time. This module generates
+  # CUDA C++ source code that uses MathDx and compiles it via NVRTC.
+  #
+  # @example Fused GEMM + epilog kernel
+  #   kernel = Ignis::MathDx::GemmKernel.new(
+  #     m: 64, n: 64, k: 64,
+  #     dtype: :float16,
+  #     epilog: :relu
+  #   )
+  #   kernel.compile!
+  #   kernel.execute(a, b, c)
+  module MathDx
+    # MathDx installation path (set via environment or default)
+    MATHDX_PATH = ENV.fetch("MATHDX_PATH", "C:/Program Files/NVIDIA/MathDx")
+
+    # Check if MathDx SDK is available
+    # @return [Boolean]
+    def self.available?
+      cublasdx_header = File.join(MATHDX_PATH, "include", "cublasdx.hpp")
+      cufftdx_header = File.join(MATHDX_PATH, "include", "cufftdx.hpp")
+      File.exist?(cublasdx_header) || File.exist?(cufftdx_header)
+    rescue StandardError
+      false
+    end
+
+    # Get MathDx include path for NVRTC compilation
+    # @return [String]
+    def self.include_path
+      File.join(MATHDX_PATH, "include")
+    end
+
+    # Supported compute capabilities for MathDx
+    SUPPORTED_SM_VERSIONS = [70, 75, 80, 86, 89, 90].freeze
+
+    class << self
+      # Create a fused GEMM kernel with device-side MathDx acceleration
+      # @param m [Integer] Output rows
+      # @param n [Integer] Output columns
+      # @param k [Integer] Inner dimension
+      # @param dtype [Symbol] Data type (:float16, :float32, :float64)
+      # @param epilog [Symbol, nil] Epilog operation (:relu, :gelu, :sigmoid, nil)
+      # @param block_size [Integer] Thread block size
+      # @return [GemmKernel]
+      def gemm_kernel(m:, n:, k:, dtype: :float32, epilog: nil, block_size: 128)
+        GemmKernel.new(m: m, n: n, k: k, dtype: dtype, epilog: epilog, block_size: block_size)
+      end
+
+      # Create a fused FFT kernel with device-side MathDx acceleration
+      # @param size [Integer] FFT size
+      # @param dtype [Symbol] Data type (:complex64, :complex128)
+      # @param direction [Symbol] Direction (:forward, :inverse)
+      # @param elements_per_thread [Integer] Elements per thread
+      # @return [FftKernel]
+      def fft_kernel(size:, dtype: :complex64, direction: :forward, elements_per_thread: 8)
+        FftKernel.new(size: size, dtype: dtype, direction: direction,
+                      elements_per_thread: elements_per_thread)
+      end
+    end
+  end
+end