ignis 0.0.1

data/lib/nvruby/linalg/matmul.rb

@@ -0,0 +1,247 @@
+# frozen_string_literal: true
+
+module Ignis
+  module LinAlg
+    # Matrix multiplication operations using cuBLAS
+    module Matmul
+      class << self
+        # Perform matrix multiplication: C = alpha * A @ B + beta * C
+        # @param a [NvArray] Left matrix
+        # @param b [NvArray] Right matrix
+        # @param c [NvArray, nil] Output matrix (created if nil)
+        # @param alpha [Float] Scaling factor for A @ B
+        # @param beta [Float] Scaling factor for C
+        # @param transpose_a [Boolean] Transpose A
+        # @param transpose_b [Boolean] Transpose B
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Result matrix
+        def call(a, b, c: nil, alpha: 1.0, beta: 0.0, transpose_a: false, transpose_b: false, stream: nil)
+          validate_inputs!(a, b)
+
+          # Get dimensions
+          m, k1, k2, n = compute_dimensions(a, b, transpose_a, transpose_b)
+          raise DimensionError, "Matrix dimensions incompatible: A(#{a.shape}) @ B(#{b.shape})" unless k1 == k2
+
+          # Ensure arrays are on device
+          a = a.to_device unless a.on_device?
+          b = b.to_device unless b.on_device?
+
+          # Create or validate output
+          c = prepare_output(c, m, n, a)
+          c = c.to_device unless c.on_device?
+
+          # Perform GEMM
+          execute_gemm(a, b, c, m, n, k1, alpha, beta, transpose_a, transpose_b, stream)
+
+          c
+        end
+
+        # Shorthand for matrix multiplication: C = A @ B
+        # @param a [NvArray] Left matrix
+        # @param b [NvArray] Right matrix
+        # @return [NvArray] Result matrix
+        def matmul(a, b)
+          call(a, b)
+        end
+
+        # Perform GEMM with a specific algorithm (useful for benchmarking/perf tuning)
+        def call_with_algo(a, b, algo, c: nil, alpha: 1.0, beta: 0.0, transpose_a: false, transpose_b: false, stream: nil)
+          validate_inputs!(a, b)
+          m, k1, k2, n = compute_dimensions(a, b, transpose_a, transpose_b)
+          raise DimensionError, "Matrix dimensions incompatible" unless k1 == k2
+
+          a = a.to_device unless a.on_device?
+          b = b.to_device unless b.on_device?
+          c = prepare_output(c, m, n, a)
+          c = c.to_device unless c.on_device?
+
+          CuBLASBindings.ensure_loaded!
+          handle = CuBLASBindings.get_handle
+
+          if stream
+            CuBLASBindings.check_status!(CuBLASBindings.cublasSetStream_v2(handle, stream.handle))
+          end
+
+          op_a = transpose_b ? CuBLASBindings::CUBLAS_OP_T : CuBLASBindings::CUBLAS_OP_N
+          op_b = transpose_a ? CuBLASBindings::CUBLAS_OP_T : CuBLASBindings::CUBLAS_OP_N
+
+          lda = b.shape[1]
+          ldb = a.shape[1]
+          ldc = n
+
+          status = execute_gemmex(handle, op_a, op_b, n, m, k1, alpha, a, b, c, lda, ldb, ldc, algo, beta)
+          CuBLASBindings.check_status!(status, "GEMM execution with algorithm #{algo}")
+          c
+        end
+
+        private
+
+        # Accept Ignis::NvArray or any array-like exposing the device-pointer API
+        # (e.g. Ignis::Shared::NvArray from the AI stack).
+        def array_like?(x)
+          # Duck-typed: any GPU array exposing the device-pointer API works
+          # (Ignis::Shared::NvArray, or Ignis::NvArray from ignis-numerics). The
+          # core no longer hard-references the numerics NvArray class.
+          x.respond_to?(:device_ffi_ptr)
+        end
+
+        # Validate input arrays
+        def validate_inputs!(a, b)
+          raise ArgumentError, "Expected NvArray, got #{a.class}" unless array_like?(a)
+          raise ArgumentError, "Expected NvArray, got #{b.class}" unless array_like?(b)
+          raise DimensionError, "Matrix A must be 2D, got #{a.ndim}D" unless a.ndim == 2
+          raise DimensionError, "Matrix B must be 2D, got #{b.ndim}D" unless b.ndim == 2
+          raise ArgumentError, "Data types must match: #{a.dtype} vs #{b.dtype}" unless a.dtype == b.dtype
+
+          dtype = a.dtype
+          unless %i[float16 bfloat16 float32 float64 complex64 complex128].include?(dtype)
+            raise UnsupportedDTypeError.new(dtype, operation: "matrix multiplication")
+          end
+        end
+
+        # GEMM algorithm constants
+        CUBLAS_GEMM_DEFAULT = -1
+        CUBLAS_GEMM_DEFAULT_TENSOR_OP = 99
+        CUBLAS_GEMM_FP16_OPTIMIZED = 1   # Peak for RTX 3060 FP16
+        CUBLAS_GEMM_FP32_OPTIMIZED = 101 # Peak for RTX 3060 FP32 TF32
+
+        # Compute dimensions considering transposition
+        # @return [Array<Integer>] [m, k1, k2, n]
+        def compute_dimensions(a, b, transpose_a, transpose_b)
+          m = transpose_a ? a.shape[1] : a.shape[0]
+          k1 = transpose_a ? a.shape[0] : a.shape[1]
+          k2 = transpose_b ? b.shape[1] : b.shape[0]
+          n = transpose_b ? b.shape[0] : b.shape[1]
+          [m, k1, k2, n]
+        end
+
+        # Prepare output array. Allocates an output of the SAME class as the
+        # input template, so the AI stack (Shared::NvArray) gets a Shared::NvArray
+        # back rather than a foreign Ignis::NvArray.
+        def prepare_output(c, m, n, template)
+          dtype = template.dtype
+          if c
+            raise DimensionError, "Output shape must be [#{m}, #{n}], got #{c.shape}" unless c.shape == [m, n]
+            raise ArgumentError, "Output dtype must be #{dtype}, got #{c.dtype}" unless c.dtype == dtype
+
+            c
+          elsif defined?(NvArray) && template.is_a?(NvArray)
+            NvArray.zeros([m, n], dtype: dtype, device: template.device_index)
+          else
+            # Array-like (e.g. Ignis::Shared::NvArray). beta defaults to 0 so cuBLAS
+            # overwrites C; we only need it allocated, not zeroed.
+            out = template.class.new(shape: [m, n], dtype: dtype, device_id: template.device_index)
+            out.to_device
+          end
+        end
+
+        # Execute GEMM operation via cuBLAS
+        # rubocop:disable Metrics/ParameterLists, Metrics/AbcSize, Metrics/MethodLength
+        def execute_gemm(a, b, c, m, n, k, alpha, beta, transpose_a, transpose_b, stream)
+          CuBLASBindings.ensure_loaded!
+          handle = CuBLASBindings.get_handle
+
+          # Ensure the handle is always associated with the correct stream.
+          # If stream is nil, we MUST reset it to 0 (default stream) because the handle
+          # might be holding a stale reference to a destroyed stream from a previous capture.
+          stream_handle = stream&.handle || FFI::Pointer::NULL
+          status = CuBLASBindings.cublasSetStream_v2(handle, stream_handle)
+          CuBLASBindings.check_status!(status, "Set cuBLAS stream")
+
+          # cuBLAS uses column-major, so we compute B^T @ A^T = (A @ B)^T
+          # by swapping A and B and transposition
+          op_a = transpose_b ? CuBLASBindings::CUBLAS_OP_T : CuBLASBindings::CUBLAS_OP_N
+          op_b = transpose_a ? CuBLASBindings::CUBLAS_OP_T : CuBLASBindings::CUBLAS_OP_N
+
+          # Leading dimensions
+          lda = b.shape[1]  # Leading dimension of B (swapped)
+          ldb = a.shape[1]  # Leading dimension of A (swapped)
+          ldc = n           # Leading dimension of C
+
+          # Prepare scalar pointers
+          case a.dtype
+          when :float32
+            # Use GemmEx for FP32 to enable TF32 and optimized algorithms
+            status = execute_gemmex(handle, op_a, op_b, n, m, k,
+                                    alpha, a, b, c, lda, ldb, ldc, CUBLAS_GEMM_FP32_OPTIMIZED, beta)
+          when :float64
+            alpha_ptr = FFI::MemoryPointer.new(:double).tap { |p| p.put_double(0, alpha) }
+            beta_ptr = FFI::MemoryPointer.new(:double).tap { |p| p.put_double(0, beta) }
+
+            status = CuBLASBindings.cublasDgemm_v2(
+              handle, op_a, op_b,
+              n, m, k,
+              alpha_ptr, b.device_ffi_ptr, lda,
+              a.device_ffi_ptr, ldb,
+              beta_ptr, c.device_ffi_ptr, ldc
+            )
+          when :complex64
+            # Complex alpha/beta need 2 floats each
+            alpha_ptr = FFI::MemoryPointer.new(:float, 2)
+            alpha_ptr.put_float(0, alpha.is_a?(Complex) ? alpha.real : alpha)
+            alpha_ptr.put_float(4, alpha.is_a?(Complex) ? alpha.imag : 0.0)
+            beta_ptr = FFI::MemoryPointer.new(:float, 2)
+            beta_ptr.put_float(0, beta.is_a?(Complex) ? beta.real : beta)
+            beta_ptr.put_float(4, beta.is_a?(Complex) ? beta.imag : 0.0)
+
+            status = CuBLASBindings.cublasCgemm_v2(
+              handle, op_a, op_b,
+              n, m, k,
+              alpha_ptr, b.device_ffi_ptr, lda,
+              a.device_ffi_ptr, ldb,
+              beta_ptr, c.device_ffi_ptr, ldc
+            )
+          when :complex128
+            alpha_ptr = FFI::MemoryPointer.new(:double, 2)
+            alpha_ptr.put_double(0, alpha.is_a?(Complex) ? alpha.real : alpha)
+            alpha_ptr.put_double(8, alpha.is_a?(Complex) ? alpha.imag : 0.0)
+            beta_ptr = FFI::MemoryPointer.new(:double, 2)
+            beta_ptr.put_double(0, beta.is_a?(Complex) ? beta.real : beta)
+            beta_ptr.put_double(8, beta.is_a?(Complex) ? beta.imag : 0.0)
+
+            status = CuBLASBindings.cublasZgemm_v2(
+              handle, op_a, op_b,
+              n, m, k,
+              alpha_ptr, b.device_ffi_ptr, lda,
+              a.device_ffi_ptr, ldb,
+              beta_ptr, c.device_ffi_ptr, ldc
+            )
+          when :float16, :bfloat16
+            # Use cublasGemmEx for half-precision with optimized algorithm
+            status = execute_gemmex(handle, op_a, op_b, n, m, k,
+                                    alpha, a, b, c, lda, ldb, ldc, CUBLAS_GEMM_FP16_OPTIMIZED, beta)
+          end
+
+          CuBLASBindings.check_status!(status, "GEMM execution")
+        end
+        # rubocop:enable Metrics/ParameterLists, Metrics/AbcSize, Metrics/MethodLength
+
+        # Execute GemmEx for optimized operations
+        def execute_gemmex(handle, op_a, op_b, m, n, k, alpha, a, b, c, lda, ldb, ldc, algo, beta = 0.0)
+          # Alpha and beta are always FP32 when using COMPUTE_32F modes
+          alpha_ptr = FFI::MemoryPointer.new(:float).tap { |p| p.put_float(0, alpha.to_f) }
+          beta_ptr = FFI::MemoryPointer.new(:float).tap { |p| p.put_float(0, beta.to_f) }
+
+          a_cuda_type = DType.cublas_type(a.dtype)
+          b_cuda_type = DType.cublas_type(b.dtype)
+          c_cuda_type = DType.cublas_type(c.dtype)
+
+          compute_type = case a.dtype
+                         when :float16 then CuBLASBindings::CUBLAS_COMPUTE_32F_FAST_16F
+                         when :bfloat16 then CuBLASBindings::CUBLAS_COMPUTE_32F_FAST_16BF
+                         when :float32 then CuBLASBindings::CUBLAS_COMPUTE_32F_FAST_TF32
+                         else CuBLASBindings::CUBLAS_COMPUTE_32F
+                         end
+
+          CuBLASBindings.cublasGemmEx(
+            handle, op_a, op_b, m, n, k,
+            alpha_ptr, b.device_ffi_ptr, b_cuda_type, lda,
+            a.device_ffi_ptr, a_cuda_type, ldb,
+            beta_ptr, c.device_ffi_ptr, c_cuda_type, ldc,
+            compute_type, algo
+          )
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/linalg/matmul_plan.rb

@@ -0,0 +1,229 @@
+# frozen_string_literal: true
+
+module Ignis
+  module LinAlg
+    # Stateful matrix multiplication plan with autotuning
+    # Reusable for repeated operations with same dimensions
+    class MatmulPlan
+      # @return [Array<Integer>] Shape of matrix A
+      attr_reader :shape_a
+
+      # @return [Array<Integer>] Shape of matrix B
+      attr_reader :shape_b
+
+      # @return [Symbol] Data type
+      attr_reader :dtype
+
+      # @return [Hash] Plan options
+      attr_reader :options
+
+      # @return [Boolean] Whether plan has been autotuned
+      attr_reader :autotuned
+
+      # @param shape_a [Array<Integer>] Shape of matrix A [m, k]
+      # @param shape_b [Array<Integer>] Shape of matrix B [k, n]
+      # @param dtype [Symbol] Data type
+      # @param transpose_a [Boolean] Transpose A
+      # @param transpose_b [Boolean] Transpose B
+      # @param epilog [Symbol, nil] Epilog operation
+      # @param device [Integer, nil] Target device
+      def initialize(shape_a:, shape_b:, dtype: :float32, transpose_a: false, transpose_b: false,
+                     epilog: nil, device: nil)
+        @shape_a = Array(shape_a)
+        @shape_b = Array(shape_b)
+        @dtype = DType.validate!(dtype)
+        @transpose_a = transpose_a
+        @transpose_b = transpose_b
+        @epilog = epilog
+        @device_index = device || Ignis.configuration.default_device
+
+        validate_shapes!
+
+        @options = {}
+        @autotuned = false
+        @best_algorithm = nil
+        @workspace = nil
+
+        @m, @k, @n = compute_dimensions
+        @execution_count = 0
+      end
+
+      # Output shape of the matmul operation
+      # @return [Array<Integer>]
+      def output_shape
+        [@m, @n]
+      end
+
+      # Plan the operation (find algorithms)
+      # @param workspace_size [Integer] Maximum workspace size in bytes
+      # @return [self]
+      def plan!(workspace_size: nil)
+        workspace_size ||= Ignis.configuration.default_workspace_size
+
+        CuBLASBindings.ensure_loaded!
+
+        Ignis.logger.debug { "Planning MatmulPlan for #{@shape_a} @ #{@shape_b} -> #{output_shape}" }
+
+        # For basic cuBLAS GEMM, planning is straightforward
+        # Advanced planning with cuBLASLt would involve algorithm selection
+        @options[:workspace_size] = workspace_size
+        @options[:planned] = true
+
+        Ignis.logger.info { "MatmulPlan planned: workspace=#{workspace_size} bytes" }
+
+        self
+      end
+
+      # Autotune the operation to find the best algorithm
+      # @param iterations [Integer] Number of benchmark iterations
+      # @param warmup [Integer] Number of warmup iterations
+      # @return [self]
+      def autotune!(iterations: nil, warmup: 3)
+        iterations ||= Ignis.configuration.autotuning_iterations
+
+        plan! unless @options[:planned]
+
+        Ignis.logger.info { "Autotuning MatmulPlan with #{iterations} iterations" }
+
+        # Create test arrays
+        a = NvArray.zeros(@shape_a, dtype: @dtype, device: @device_index)
+        b = NvArray.zeros(@shape_b, dtype: @dtype, device: @device_index)
+        c = NvArray.zeros(output_shape, dtype: @dtype, device: @device_index)
+
+        # Warmup
+        warmup.times { execute_internal(a, b, c) }
+        CUDA::Device.current.synchronize
+
+        # Benchmark
+        start_event = CUDA::Event.new
+        end_event = CUDA::Event.new
+
+        start_event.record
+        iterations.times { execute_internal(a, b, c) }
+        end_event.record
+        end_event.synchronize
+
+        elapsed_ms = CUDA::Event.elapsed_time(start_event, end_event)
+        avg_time = elapsed_ms / iterations
+
+        @options[:avg_time_ms] = avg_time
+        @autotuned = true
+
+        # Cleanup
+        start_event.destroy!
+        end_event.destroy!
+        a.free!
+        b.free!
+        c.free!
+
+        Ignis.logger.info { "MatmulPlan autotuned: avg_time=#{avg_time.round(3)}ms" }
+
+        self
+      end
+
+      # Execute the planned matrix multiplication
+      # @param a [NvArray] Left matrix
+      # @param b [NvArray] Right matrix
+      # @param c [NvArray, nil] Output matrix (created if nil)
+      # @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
+      def execute(a, b, c: nil, alpha: 1.0, beta: 0.0, stream: nil)
+        validate_execution_inputs!(a, b)
+
+        # Ensure on device
+        a = a.to_device(device: @device_index) unless a.on_device?
+        b = b.to_device(device: @device_index) unless b.on_device?
+
+        # Prepare output
+        if c
+          validate_output!(c)
+          c = c.to_device(device: @device_index) unless c.on_device?
+        else
+          c = NvArray.zeros(output_shape, dtype: @dtype, device: @device_index)
+        end
+
+        execute_internal(a, b, c, alpha, beta, stream)
+
+        c
+      end
+
+      # Get statistics about the plan
+      # @return [Hash]
+      def stats
+        {
+          shape_a: @shape_a,
+          shape_b: @shape_b,
+          output_shape: output_shape,
+          dtype: @dtype,
+          transpose_a: @transpose_a,
+          transpose_b: @transpose_b,
+          autotuned: @autotuned,
+          avg_time_ms: @options[:avg_time_ms],
+          execution_count: @execution_count
+        }
+      end
+
+      # @return [String]
+      def to_s
+        tuned = @autotuned ? "autotuned" : "not tuned"
+        "MatmulPlan(#{@shape_a} @ #{@shape_b} -> #{output_shape}, #{@dtype}, #{tuned})"
+      end
+
+      private
+
+      # Validate input shapes
+      def validate_shapes!
+        raise DimensionError, "Matrix A must be 2D" unless @shape_a.size == 2
+        raise DimensionError, "Matrix B must be 2D" unless @shape_b.size == 2
+
+        # Check K dimension compatibility
+        k_a = @transpose_a ? @shape_a[0] : @shape_a[1]
+        k_b = @transpose_b ? @shape_b[1] : @shape_b[0]
+
+        raise DimensionError, "Matrix dimensions incompatible: #{@shape_a} @ #{@shape_b}" unless k_a == k_b
+      end
+
+      # Compute output dimensions
+      # @return [Array<Integer>] [m, k, n]
+      def compute_dimensions
+        m = @transpose_a ? @shape_a[1] : @shape_a[0]
+        k = @transpose_a ? @shape_a[0] : @shape_a[1]
+        n = @transpose_b ? @shape_b[0] : @shape_b[1]
+        [m, k, n]
+      end
+
+      # Validate execution inputs
+      def validate_execution_inputs!(a, b)
+        raise ArgumentError, "Expected NvArray, got #{a.class}" unless a.is_a?(NvArray)
+        raise ArgumentError, "Expected NvArray, got #{b.class}" unless b.is_a?(NvArray)
+        raise DimensionError, "A shape mismatch: expected #{@shape_a}, got #{a.shape}" unless a.shape == @shape_a
+        raise DimensionError, "B shape mismatch: expected #{@shape_b}, got #{b.shape}" unless b.shape == @shape_b
+        raise ArgumentError, "A dtype mismatch: expected #{@dtype}, got #{a.dtype}" unless a.dtype == @dtype
+        raise ArgumentError, "B dtype mismatch: expected #{@dtype}, got #{b.dtype}" unless b.dtype == @dtype
+      end
+
+      # Validate output array
+      def validate_output!(c)
+        expected = output_shape
+        raise DimensionError, "Output shape mismatch: expected #{expected}, got #{c.shape}" unless c.shape == expected
+        raise ArgumentError, "Output dtype mismatch: expected #{@dtype}, got #{c.dtype}" unless c.dtype == @dtype
+      end
+
+      # Execute GEMM internally
+      def execute_internal(a, b, c, alpha = 1.0, beta = 0.0, stream = nil)
+        @execution_count += 1
+        Matmul.call(
+          a, b,
+          c: c,
+          alpha: alpha,
+          beta: beta,
+          transpose_a: @transpose_a,
+          transpose_b: @transpose_b,
+          stream: stream
+        )
+      end
+    end
+  end
+end