ignis-numerics 0.0.1

data/lib/nvruby/array.rb

@@ -0,0 +1,646 @@
+# frozen_string_literal: true
+
+module Ignis
+  # GPU-aware multi-dimensional array
+  # Similar to NumPy ndarray but with CUDA memory backing
+  class NvArray
+    # @return [Array<Integer>] Shape of the array
+    attr_reader :shape
+
+    # @return [Symbol] Data type (:float32, :float64, :complex64, etc.)
+    attr_reader :dtype
+
+    # @return [Array<Integer>] Strides in bytes for each dimension
+    attr_reader :strides
+
+    # @return [CUDA::Memory, nil] Device memory (nil if on host)
+    attr_reader :device_memory
+
+    # @return [FFI::Pointer, nil] Host memory pointer (nil if on device)
+    attr_reader :host_memory
+
+    # @return [:host, :device] Current memory location
+    attr_reader :location
+
+    # @return [Integer] Device index
+    attr_reader :device_index
+
+    # Create a new NvArray
+    # @param shape [Array<Integer>] Shape of the array
+    # @param dtype [Symbol] Data type (default: :float32)
+    # @param device [Integer, nil] Device index (nil for host-only)
+    # @param data [Array, String, nil] Initial data
+    def initialize(shape:, dtype: :float32, device: nil, data: nil)
+      @shape = normalize_shape(shape)
+      @dtype = DType.validate!(dtype)
+      @device_index = device || Ignis.configuration.default_device
+      @strides = compute_strides
+      @size_bytes = compute_size_bytes
+
+      @device_memory = nil
+      @host_memory = nil
+
+      if data
+        initialize_with_data(data)
+      else
+        allocate_memory(device ? :device : :host)
+      end
+    end
+
+    # @return [Integer] Total number of elements
+    def size
+      @shape.reduce(1, :*)
+    end
+
+    # @return [Integer] Number of dimensions
+    def ndim
+      @shape.size
+    end
+
+    # @return [Integer] Total size in bytes
+    def nbytes
+      @size_bytes
+    end
+
+    # @return [Integer] Size of each element in bytes
+    def itemsize
+      DType.byte_size(@dtype)
+    end
+
+    # Check if data is on device
+    # @return [Boolean]
+    def on_device?
+      @location == :device
+    end
+
+    # Check if data is on host
+    # @return [Boolean]
+    def on_host?
+      @location == :host
+    end
+
+    # Transfer data to GPU
+    # @param device [Integer, nil] Target device (nil for current)
+    # @param stream [CUDA::Stream, nil] Stream for async transfer
+    # @return [self]
+    def to_device(device: nil, stream: nil)
+      return self if on_device? && (device.nil? || device == @device_index)
+
+      target_device = device || @device_index
+
+      ensure_host_data!
+
+      @device_memory = CUDA::Memory.new(@size_bytes, device: target_device)
+      @device_memory.copy_from_host(@host_memory, stream: stream)
+
+      @device_index = target_device
+      @location = :device
+
+      # Free host memory if not needed
+      @host_memory = nil unless Ignis.configuration.with_lock { false } # Keep host copy option
+
+      self
+    end
+
+    # Transfer data to host
+    # @param stream [CUDA::Stream, nil] Stream for async transfer
+    # @return [self]
+    def to_host(stream: nil)
+      return self if on_host?
+
+      ensure_device_data!
+
+      @host_memory = FFI::MemoryPointer.new(:uint8, @size_bytes)
+      @device_memory.copy_to_host(host_buffer: @host_memory, stream: stream)
+
+      @location = :host
+      @device_memory.free!
+      @device_memory = nil
+
+      self
+    end
+
+    # Get device pointer for CUDA operations
+    # @return [Fiddle::Pointer] Device pointer
+    # @raise [InvalidOperationError] If not on device
+    def device_ptr
+      ensure_device_data!
+      @device_memory.device_ptr
+    end
+
+    # Get device pointer wrapped as an FFI::Pointer for FFI-bound CUDA-X
+    # library calls (cuBLAS/cuSOLVER/cuFFT/cuRAND/cuSPARSE), which cannot
+    # accept the Fiddle::Pointer returned by #device_ptr.
+    # @return [FFI::Pointer] Device pointer
+    # @raise [InvalidOperationError] If not on device
+    def device_ffi_ptr
+      ensure_device_data!
+      @device_memory.ffi_ptr
+    end
+
+    # Get host pointer
+    # @return [FFI::Pointer] Host pointer
+    # @raise [InvalidOperationError] If not on host
+    def host_ptr
+      ensure_host_data!
+      @host_memory
+    end
+
+    # Get data as Ruby array (copies to host if needed)
+    # @return [Array] Nested Ruby array with data
+    def to_a
+      synchronize_if_needed
+      ensure_host_data!
+
+      flat_data = read_flat_data
+      reshape_to_nested(flat_data, @shape)
+    end
+
+    # Get flat data as Ruby array
+    # @return [Array] 1D Ruby array with all elements
+    def flatten
+      synchronize_if_needed
+      ensure_host_data!
+      read_flat_data
+    end
+
+    # Reshape the array
+    # @param new_shape [Array<Integer>] New shape
+    # @return [NvArray] Reshaped array (view if contiguous)
+    def reshape(new_shape)
+      new_shape = normalize_shape(new_shape)
+
+      # Handle -1 in shape
+      if new_shape.include?(-1)
+        neg_idx = new_shape.index(-1)
+        other_size = new_shape.reject { |d| d == -1 }.reduce(1, :*)
+        new_shape[neg_idx] = size / other_size
+      end
+
+      raise DimensionError, "Cannot reshape array of size #{size} to #{new_shape}" unless new_shape.reduce(1, :*) == size
+
+      # Create new array with same memory
+      result = dup
+      result.instance_variable_set(:@shape, new_shape)
+      result.instance_variable_set(:@strides, result.send(:compute_strides))
+      result
+    end
+
+    # Transpose the array
+    # @param axes [Array<Integer>, nil] Permutation of axes (reverses if nil)
+    # @return [NvArray] Transposed array
+    def transpose(axes: nil)
+      axes ||= (0...ndim).to_a.reverse
+
+      raise DimensionError, "Invalid axes for transpose" unless axes.sort == (0...ndim).to_a
+
+      new_shape = axes.map { |ax| @shape[ax] }
+      new_strides = axes.map { |ax| @strides[ax] }
+
+      result = dup
+      result.instance_variable_set(:@shape, new_shape)
+      result.instance_variable_set(:@strides, new_strides)
+      result
+    end
+
+    # Create a contiguous copy
+    # @return [NvArray] Contiguous copy
+    def contiguous
+      return self if contiguous?
+
+      # Create new array and copy data
+      result = NvArray.new(shape: @shape, dtype: @dtype, device: on_device? ? @device_index : nil)
+
+      if on_device?
+        # Device-to-device copy
+        result.device_memory.copy_from_device(@device_memory)
+      else
+        # Copy host data
+        result.host_memory.put_bytes(0, @host_memory.get_bytes(0, @size_bytes))
+      end
+
+      result
+    end
+
+    # Check if memory layout is contiguous
+    # @return [Boolean]
+    def contiguous?
+      expected = itemsize
+      @strides.reverse.each_with_index do |stride, i|
+        dim = @shape[ndim - 1 - i]
+        return false unless stride == expected || dim == 1
+
+        expected *= dim
+      end
+      true
+    end
+
+    # Duplicate the array
+    # @return [NvArray] Copy of the array
+    def dup
+      result = NvArray.new(shape: @shape.dup, dtype: @dtype, device: on_device? ? @device_index : nil)
+
+      if on_device?
+        result.instance_variable_set(:@device_memory, CUDA::Memory.new(@size_bytes, device: @device_index))
+        result.device_memory.copy_from_device(@device_memory)
+        result.instance_variable_set(:@location, :device)
+      else
+        new_host = FFI::MemoryPointer.new(:uint8, @size_bytes)
+        new_host.put_bytes(0, @host_memory.get_bytes(0, @size_bytes))
+        result.instance_variable_set(:@host_memory, new_host)
+        result.instance_variable_set(:@location, :host)
+      end
+
+      result
+    end
+
+    # Zero out the array
+    # @param stream [CUDA::Stream, nil] Stream for async operation
+    # @return [self]
+    def zero!(stream: nil)
+      if on_device?
+        @device_memory.zero!(stream: stream)
+      else
+        @host_memory.clear
+      end
+      self
+    end
+
+    # Free all memory
+    # @return [void]
+    def free!
+      @device_memory&.free!
+      @device_memory = nil
+      @host_memory = nil
+      @location = nil
+    end
+
+    # @return [String] String representation
+    def to_s
+      loc = on_device? ? "device:#{@device_index}" : "host"
+      "NvArray(shape=#{@shape}, dtype=#{@dtype}, #{loc})"
+    end
+
+    # @return [String] Detailed inspection
+    def inspect
+      "#<Ignis::NvArray:#{object_id} shape=#{@shape} dtype=#{@dtype} " \
+        "location=#{@location} device=#{@device_index} bytes=#{@size_bytes}>"
+    end
+
+    class << self
+      # Create array filled with zeros
+      # @param shape [Array<Integer>] Shape
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def zeros(shape, dtype: :float32, device: nil)
+        arr = new(shape: shape, dtype: dtype, device: device)
+        arr.zero!
+        arr
+      end
+
+      # Create array filled with ones
+      # @param shape [Array<Integer>] Shape
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def ones(shape, dtype: :float32, device: nil)
+        size = Array(shape).reduce(1, :*)
+        data = Array.new(size, 1.0)
+        new(shape: shape, dtype: dtype, device: device, data: data)
+      end
+
+      # Create array with evenly spaced values
+      # @param start [Numeric] Start value
+      # @param stop [Numeric] End value
+      # @param num [Integer] Number of samples
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def linspace(start, stop, num, dtype: :float32, device: nil)
+        step = (stop - start).to_f / (num - 1)
+        data = (0...num).map { |i| start + step * i }
+        new(shape: [num], dtype: dtype, device: device, data: data)
+      end
+
+      # Create array from existing device memory
+      # @param ptr [FFI::Pointer, CUDA::Memory] Device pointer or Memory object
+      # @param shape [Array<Integer>] Shape
+      # @param dtype [Symbol] Data type
+      # @param take_ownership [Boolean] If true, the array will manage the memory lifecycle
+      # @return [NvArray]
+      def from_device_ptr(ptr, shape:, dtype: :float32, take_ownership: false)
+        arr = allocate_empty_metadata(shape, dtype)
+        
+        if ptr.is_a?(CUDA::Memory)
+          # If it's already a Memory object, we can use it directly
+          # If we don't take ownership, we might need a non-owning version
+          if take_ownership
+            arr.instance_variable_set(:@device_memory, ptr)
+          else
+            # Create a non-owning wrapper for the same pointer
+            wrapper = CUDA::Memory.new(arr.nbytes, device: ptr.device_index, ptr: ptr.device_ptr, owned: false)
+            arr.instance_variable_set(:@device_memory, wrapper)
+          end
+        else
+          # It's a raw FFI::Pointer
+          # We wrap it in a Memory object
+          wrapper = CUDA::Memory.new(arr.nbytes, ptr: ptr, owned: take_ownership)
+          arr.instance_variable_set(:@device_memory, wrapper)
+        end
+        
+        arr.instance_variable_set(:@location, :device)
+        arr.instance_variable_set(:@device_index, arr.device_memory.device_index)
+        arr
+      end
+
+      # Create an empty NvArray without allocating memory
+      # @param shape [Array<Integer>] Shape
+      # @param dtype [Symbol] Data type
+      # @return [NvArray]
+      def allocate_empty_metadata(shape, dtype)
+        # We use allocate to avoid calling initialize which tries to allocate memory
+        arr = allocate
+        arr.instance_variable_set(:@shape, Array(shape).map(&:to_i))
+        arr.instance_variable_set(:@dtype, DType.validate!(dtype))
+        arr.instance_variable_set(:@strides, arr.send(:compute_strides))
+        arr.instance_variable_set(:@size_bytes, arr.send(:compute_size_bytes))
+        arr
+      end
+
+      # Create array filled with ones
+      # @param shape [Array<Integer>] Shape
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def ones(shape, dtype: :float32, device: nil)
+        size = Array(shape).reduce(1, :*)
+        data = Array.new(size, 1.0)
+        new(shape: shape, dtype: dtype, device: device, data: data)
+      end
+
+      # Create array with evenly spaced values
+      # @param start [Numeric] Start value
+      # @param stop [Numeric] End value
+      # @param num [Integer] Number of samples
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def linspace(start, stop, num, dtype: :float32, device: nil)
+        step = (stop - start).to_f / (num - 1)
+        data = (0...num).map { |i| start + step * i }
+        new(shape: [num], dtype: dtype, device: device, data: data)
+      end
+
+      # Create array from Ruby array
+      # @param data [Array] Nested Ruby array
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def from_array(data, dtype: :float32, device: nil)
+        shape = infer_shape(data)
+        flat = flatten_nested(data)
+        new(shape: shape, dtype: dtype, device: device, data: flat)
+      end
+
+      # Create identity matrix
+      # @param size [Integer] Size of the square matrix
+      # @param dtype [Symbol] Data type
+      # @param device [Integer, nil] Device index
+      # @return [NvArray]
+      def eye(size, dtype: :float32, device: nil)
+        data = Array.new(size * size) { |i| i / size == i % size ? 1.0 : 0.0 }
+        new(shape: [size, size], dtype: dtype, device: device, data: data)
+      end
+
+      private
+
+      # Infer shape from nested array
+      # @param data [Array] Nested array
+      # @return [Array<Integer>]
+      def infer_shape(data)
+        shape = []
+        current = data
+        while current.is_a?(Array)
+          shape << current.size
+          current = current.first
+        end
+        shape
+      end
+
+      # Flatten nested array
+      # @param data [Array] Nested array
+      # @return [Array]
+      def flatten_nested(data)
+        data.flatten
+      end
+    end
+
+    private
+
+    # Normalize shape input
+    # @param shape [Array<Integer>, Integer]
+    # @return [Array<Integer>]
+    def normalize_shape(shape)
+      Array(shape).map(&:to_i)
+    end
+
+    # Compute strides for row-major layout
+    # @return [Array<Integer>]
+    def compute_strides
+      strides = []
+      stride = itemsize
+      @shape.reverse_each do |dim|
+        strides.unshift(stride)
+        stride *= dim
+      end
+      strides
+    end
+
+    # Compute total size in bytes
+    # @return [Integer]
+    def compute_size_bytes
+      size * itemsize
+    end
+
+    # Allocate memory on specified location
+    # @param location [:host, :device]
+    def allocate_memory(location)
+      @location = location
+
+      if location == :device
+        @device_memory = CUDA::Memory.new(@size_bytes, device: @device_index)
+      else
+        @host_memory = FFI::MemoryPointer.new(:uint8, @size_bytes)
+      end
+    end
+
+    # Initialize with provided data
+    # @param data [Array, String]
+    def initialize_with_data(data)
+      @host_memory = FFI::MemoryPointer.new(:uint8, @size_bytes)
+      @location = :host
+
+      flat_data = data.is_a?(Array) ? data.flatten : data
+
+      write_data(flat_data)
+    end
+
+    # Write data to host memory
+    # @param data [Array, String]
+    def write_data(data)
+      case data
+      when String
+        @host_memory.put_bytes(0, data)
+      when Array
+        if DType.complex?(@dtype)
+          write_complex_data(data)
+        else
+          write_scalar_data(data)
+        end
+      end
+    end
+
+    # Write complex data (stored as [real0, imag0, real1, imag1, ...])
+    # @param data [Array] Interleaved real/imag values
+    def write_complex_data(data)
+      component_type = DType.real_dtype(@dtype)
+      component_size = DType.byte_size(component_type)
+
+      data.each_with_index do |val, i|
+        offset = i * component_size
+        case component_type
+        when :float32
+          @host_memory.put_float(offset, val.to_f)
+        when :float64
+          @host_memory.put_double(offset, val.to_f)
+        end
+      end
+    end
+
+    # Write scalar (non-complex) or component data
+    # @param data [Array] Values to write
+    def write_scalar_data(data)
+      # Half-precision types have no native FFI representation: encode the float
+      # value into its 16-bit bit pattern rather than truncating it to an integer.
+      if @dtype == :float16
+        data.each_with_index { |val, i| @host_memory.put_uint16(i * 2, Half.f32_to_f16(val)) }
+        return
+      elsif @dtype == :bfloat16
+        data.each_with_index { |val, i| @host_memory.put_uint16(i * 2, Half.f32_to_bf16(val)) }
+        return
+      end
+
+      ffi_type = DType.ffi_type(@dtype)
+      comp_size = DType.complex?(@dtype) ? itemsize / 2 : itemsize
+      data.each_with_index do |val, i|
+        offset = i * comp_size
+        case ffi_type
+        when :float
+          @host_memory.put_float(offset, val.to_f)
+        when :double
+          @host_memory.put_double(offset, val.to_f)
+        when :int8
+          @host_memory.put_int8(offset, val.to_i)
+        when :int16
+          @host_memory.put_int16(offset, val.to_i)
+        when :int32
+          @host_memory.put_int32(offset, val.to_i)
+        when :int64
+          @host_memory.put_int64(offset, val.to_i)
+        when :uint8
+          @host_memory.put_uint8(offset, val.to_i)
+        when :uint16
+          @host_memory.put_uint16(offset, val.to_i)
+        when :uint32
+          @host_memory.put_uint32(offset, val.to_i)
+        when :uint64
+          @host_memory.put_uint64(offset, val.to_i)
+        end
+      end
+    end
+
+    # Read flat data from host memory (logically ordered)
+    # @return [Array]
+    def read_flat_data
+      ensure_host_data!
+      
+      ffi_type = DType.ffi_type(@dtype)
+      is_complex = DType.complex?(@dtype)
+      count = is_complex ? size * 2 : size
+      comp_size = is_complex ? itemsize / 2 : itemsize
+
+      # Half-precision types are stored as raw uint16 bits and must be decoded.
+      half_decoder = case @dtype
+                     when :float16 then Half.method(:f16_to_f32)
+                     when :bfloat16 then Half.method(:bf16_to_f32)
+                     end
+
+      flat = if contiguous?
+               count.times.map do |i|
+                 @host_memory.send("get_#{ffi_type}", i * comp_size)
+               end
+             else
+               # Logical iteration for non-contiguous arrays
+               values = []
+               iterate_shape do |offset|
+                 if is_complex
+                   values << @host_memory.send("get_#{ffi_type}", offset)
+                   values << @host_memory.send("get_#{ffi_type}", offset + comp_size)
+                 else
+                   values << @host_memory.send("get_#{ffi_type}", offset)
+                 end
+               end
+               values
+             end
+
+      half_decoder ? flat.map { |bits| half_decoder.call(bits) } : flat
+    end
+
+    # Iterate over shape and yield byte offsets
+    def iterate_shape(dim = 0, current_offset = 0, &block)
+      if dim == ndim - 1
+        @shape[dim].times do |i|
+          yield current_offset + i * @strides[dim]
+        end
+      else
+        @shape[dim].times do |i|
+          iterate_shape(dim + 1, current_offset + i * @strides[dim], &block)
+        end
+      end
+    end
+
+    # Reshape flat array to nested structure
+    # @param flat [Array] Flat array
+    # @param shape [Array<Integer>] Target shape
+    # @return [Array]
+    def reshape_to_nested(flat, shape)
+      return flat.first if shape.empty?
+      return flat if shape.size == 1
+
+      chunk_size = shape[1..].reduce(1, :*)
+      flat.each_slice(chunk_size).map do |chunk|
+        reshape_to_nested(chunk, shape[1..])
+      end
+    end
+
+    # Ensure data is on host
+    def ensure_host_data!
+      to_host if on_device?
+      raise InvalidOperationError, "No host data available" unless @host_memory
+    end
+
+    # Ensure data is on device
+    def ensure_device_data!
+      to_device if on_host?
+      raise InvalidOperationError, "No device data available" unless @device_memory
+    end
+
+    # Synchronize if on device
+    def synchronize_if_needed
+      CUDA::Device.current.synchronize if on_device?
+    rescue StandardError
+      # Ignore sync errors when reading
+    end
+  end
+end