ignis 0.0.1

data/lib/nvruby/cuda/memory.rb

@@ -0,0 +1,410 @@
+# frozen_string_literal: true
+
+require 'fiddle'
+
+module Ignis
+  module CUDA
+    # Manages GPU device memory allocation and transfers.
+    #
+    # Refactored to use Fiddle-based RuntimeAPI hot-path methods.
+    # No FFI::MemoryPointer usage — all pointers via Fiddle::Pointer.
+    class Memory
+      # @return [Fiddle::Pointer] Device pointer
+      attr_reader :device_ptr
+
+      # @return [Integer] Size in bytes
+      attr_reader :size
+
+      # @return [Integer] Device index
+      attr_reader :device_index
+
+      # @param size [Integer] Size in bytes to allocate
+      # @param device [Device, Integer, nil] Device to allocate on (nil for current)
+      # @param ptr [Fiddle::Pointer, nil] Existing pointer to wrap (nil to allocate new)
+      # @param owned [Boolean] Whether this object owns the memory and should free it
+      def initialize(size, device: nil, ptr: nil, owned: true)
+        @size = size
+        @device_index = resolve_device_index(device)
+        @owned = owned
+
+        if ptr
+          @device_ptr = ptr.is_a?(Fiddle::Pointer) ? ptr : Fiddle::Pointer.new(ptr.to_i)
+        else
+          @device_ptr = allocate_device_memory
+        end
+
+        @freed = false
+
+        if @owned
+          captured_ptr = @device_ptr
+          captured_device = @device_index
+          ObjectSpace.define_finalizer(self, self.class.release_finalizer(captured_ptr, captured_device))
+        end
+      end
+
+      # @return [Fiddle::Pointer] pointer for interop
+      def to_ptr
+        @device_ptr
+      end
+
+      # Device pointer wrapped as an FFI::Pointer.
+      #
+      # The hot path (this class' own cudaMemcpy etc.) is Fiddle-based, but the
+      # CUDA-X library bindings (cuBLAS/cuSOLVER/cuFFT/cuRAND/cuSPARSE) are FFI
+      # and cannot accept a Fiddle::Pointer. Use this when handing a device
+      # buffer to an FFI-bound library call.
+      # @return [FFI::Pointer]
+      def ffi_ptr
+        FFI::Pointer.new(@device_ptr.to_i)
+      end
+
+      # @return [Integer] raw address
+      def address
+        @device_ptr.to_i
+      end
+
+      # @return [Boolean]
+      def freed?
+        @freed
+      end
+
+      # Free the device memory.
+      # @return [void]
+      def free!
+        return if @freed
+        return unless @owned
+
+        RuntimeAPI.ensure_loaded!
+        ensure_correct_device do
+          RuntimeAPI.free(@device_ptr)
+        end
+
+        @freed = true
+        ObjectSpace.undefine_finalizer(self)
+      end
+
+      # Copy data from host to device.
+      # @param host_data [Fiddle::Pointer, String] Source data
+      # @param count [Integer, nil] Number of bytes to copy (defaults to size)
+      # @param stream [Stream, nil] Optional stream for async copy
+      # @return [void]
+      def copy_from_host(host_data, count: nil, stream: nil)
+        raise MemoryError, 'Memory has been freed' if @freed
+
+        count ||= @size
+        raise MemoryError, "Copy count #{count} exceeds allocation size #{@size}" if count > @size
+
+        RuntimeAPI.ensure_loaded!
+        host_ptr = prepare_host_pointer(host_data)
+
+        ensure_correct_device do
+          if stream
+            RuntimeAPI.memcpy_async(
+              @device_ptr, host_ptr, count,
+              RuntimeAPI::MEMCPY_HOST_TO_DEVICE, stream.to_ptr
+            )
+          else
+            RuntimeAPI.memcpy(
+              @device_ptr, host_ptr, count,
+              RuntimeAPI::MEMCPY_HOST_TO_DEVICE
+            )
+          end
+        end
+      end
+
+      # Copy data from device to host.
+      # @param host_buffer [Fiddle::Pointer, nil] Destination buffer (created if nil)
+      # @param count [Integer, nil] Number of bytes to copy (defaults to size)
+      # @param stream [Stream, nil] Optional stream for async copy
+      # @return [Fiddle::Pointer] Host buffer with data
+      def copy_to_host(host_buffer: nil, count: nil, stream: nil)
+        raise MemoryError, 'Memory has been freed' if @freed
+
+        count ||= @size
+        raise MemoryError, "Copy count #{count} exceeds allocation size #{@size}" if count > @size
+
+        RuntimeAPI.ensure_loaded!
+        host_buffer ||= Fiddle::Pointer.malloc(count)
+
+        # The destination may be an FFI::MemoryPointer (NvArray host buffer);
+        # bridge it to a Fiddle address for the Fiddle-based memcpy, but return
+        # the original object so the caller can read it back with its own API.
+        dst_ptr = if host_buffer.is_a?(Fiddle::Pointer)
+                    host_buffer
+                  elsif host_buffer.respond_to?(:address)
+                    Fiddle::Pointer.new(host_buffer.address)
+                  else
+                    host_buffer
+                  end
+
+        ensure_correct_device do
+          if stream
+            RuntimeAPI.memcpy_async(
+              dst_ptr, @device_ptr, count,
+              RuntimeAPI::MEMCPY_DEVICE_TO_HOST, stream.to_ptr
+            )
+          else
+            RuntimeAPI.memcpy(
+              dst_ptr, @device_ptr, count,
+              RuntimeAPI::MEMCPY_DEVICE_TO_HOST
+            )
+          end
+        end
+
+        host_buffer
+      end
+
+      # Copy data from another device memory.
+      # @param source [Memory] Source device memory
+      # @param count [Integer, nil] Number of bytes to copy
+      # @param stream [Stream, nil] Optional stream for async copy
+      # @return [void]
+      def copy_from_device(source, count: nil, stream: nil)
+        raise MemoryError, 'Memory has been freed' if @freed
+        raise MemoryError, 'Source memory has been freed' if source.freed?
+
+        count ||= [source.size, @size].min
+        raise MemoryError, "Copy count #{count} exceeds allocation size #{@size}" if count > @size
+
+        RuntimeAPI.ensure_loaded!
+
+        if stream
+          RuntimeAPI.memcpy_async(
+            @device_ptr, source.device_ptr, count,
+            RuntimeAPI::MEMCPY_DEVICE_TO_DEVICE, stream.to_ptr
+          )
+        else
+          RuntimeAPI.memcpy(
+            @device_ptr, source.device_ptr, count,
+            RuntimeAPI::MEMCPY_DEVICE_TO_DEVICE
+          )
+        end
+      end
+
+      # Set memory to a value.
+      # @param value [Integer] Byte value to set (0-255)
+      # @param count [Integer, nil] Number of bytes to set (defaults to size)
+      # @param stream [Stream, nil] Optional stream for async operation
+      # @return [void]
+      def memset(value, count: nil, stream: nil)
+        raise MemoryError, 'Memory has been freed' if @freed
+
+        count ||= @size
+        raise MemoryError, "Memset count #{count} exceeds allocation size #{@size}" if count > @size
+
+        RuntimeAPI.ensure_loaded!
+
+        ensure_correct_device do
+          if stream
+            RuntimeAPI.memset_async(@device_ptr, value, count, stream.to_ptr)
+          else
+            RuntimeAPI.memset(@device_ptr, value, count)
+          end
+        end
+      end
+
+      # Zero out the memory.
+      # @param stream [Stream, nil] Optional stream
+      # @return [void]
+      def zero!(stream: nil)
+        memset(0, stream: stream)
+      end
+
+      # @return [String]
+      def to_s
+        status = @freed ? 'freed' : 'allocated'
+        "DeviceMemory[#{@size} bytes, device #{@device_index}, #{status}]"
+      end
+
+      # @return [String]
+      def inspect
+        "#<Ignis::CUDA::Memory:#{object_id} size=#{@size} device=#{@device_index} " \
+          "ptr=0x#{@device_ptr.to_i.to_s(16)} freed=#{@freed}>"
+      end
+
+      class << self
+        # Create a finalizer proc for releasing device memory.
+        # @param ptr [Fiddle::Pointer] Device pointer to free
+        # @param device_index [Integer] Device the memory is on
+        # @return [Proc]
+        def release_finalizer(ptr, device_index)
+          ptr_addr = ptr.to_i
+          proc do
+            begin
+              RuntimeAPI.ensure_loaded!
+              current = RuntimeAPI.get_device
+              RuntimeAPI.set_device(device_index) if current != device_index
+              RuntimeAPI.free(Fiddle::Pointer.new(ptr_addr))
+              RuntimeAPI.set_device(current) if current != device_index
+            rescue StandardError
+              # Silently ignore errors during finalization
+            end
+          end
+        end
+      end
+
+      private
+
+      # Resolve device index from various inputs.
+      # @param device [Device, Integer, nil]
+      # @return [Integer]
+      def resolve_device_index(device)
+        case device
+        when Device
+          device.index
+        when Integer
+          device
+        when nil
+          Ignis.configuration.default_device
+        else
+          raise ArgumentError, "Invalid device: #{device.inspect}"
+        end
+      end
+
+      # Allocate device memory via RuntimeAPI (Fiddle hot path).
+      # Restores the previously-current device afterward so that allocating on a
+      # non-default GPU doesn't silently change the process-wide current device
+      # (which would misdirect later default-device kernels/copies on multi-GPU).
+      # @return [Fiddle::Pointer]
+      def allocate_device_memory
+        RuntimeAPI.ensure_loaded!
+        original = RuntimeAPI.get_device
+        RuntimeAPI.set_device(@device_index) if original != @device_index
+        ptr = RuntimeAPI.malloc(@size)
+        RuntimeAPI.set_device(original) if original != @device_index
+        ptr
+      end
+
+      # Ensure operations run on the correct device.
+      # @yield Block to run with correct device set
+      # @return [Object] Block return value
+      def ensure_correct_device
+        RuntimeAPI.ensure_loaded!
+        original = RuntimeAPI.get_device
+
+        RuntimeAPI.set_device(@device_index) if original != @device_index
+
+        result = yield
+
+        RuntimeAPI.set_device(original) if original != @device_index
+
+        result
+      end
+
+      # Prepare host pointer from various inputs.
+      #
+      # NvArray stores its host buffer as an FFI::MemoryPointer, while this class
+      # is Fiddle-based, so we bridge any FFI pointer (or other address-exposing
+      # object) into a non-owning Fiddle::Pointer at the same address. The caller
+      # must keep the source object alive for the duration of the (synchronous)
+      # copy — which NvArray does, since it holds the reference.
+      # @param data [Fiddle::Pointer, FFI::Pointer, String]
+      # @return [Fiddle::Pointer]
+      def prepare_host_pointer(data)
+        case data
+        when Fiddle::Pointer
+          data
+        when String
+          ptr = Fiddle::Pointer.malloc(data.bytesize)
+          ptr[0, data.bytesize] = data
+          ptr
+        else
+          raise ArgumentError, "Unsupported host data type: #{data.class}" unless data.respond_to?(:address)
+
+          Fiddle::Pointer.new(data.address)
+        end
+      end
+    end
+
+    # Pinned (page-locked) host memory for faster transfers.
+    #
+    # Uses RuntimeAPI.host_alloc (Fiddle) instead of FFI::MemoryPointer.
+    class PinnedMemory
+      # @return [Fiddle::Pointer] Host pointer
+      attr_reader :host_ptr
+
+      # @return [Integer] Size in bytes
+      attr_reader :size
+
+      # @param size [Integer] Size in bytes to allocate
+      def initialize(size)
+        @size = size
+        @host_ptr = allocate_pinned_memory
+        @freed = false
+
+        captured_ptr = @host_ptr
+        ObjectSpace.define_finalizer(self, self.class.release_finalizer(captured_ptr))
+      end
+
+      # @return [Boolean]
+      def freed?
+        @freed
+      end
+
+      # Free the pinned memory.
+      # @return [void]
+      def free!
+        return if @freed
+
+        RuntimeAPI.ensure_loaded!
+        RuntimeAPI.free_host(@host_ptr)
+
+        @freed = true
+        ObjectSpace.undefine_finalizer(self)
+      end
+
+      # Write data to the pinned memory.
+      # @param data [String] Data to write
+      # @param offset [Integer] Offset in bytes
+      # @return [void]
+      def write(data, offset: 0)
+        raise MemoryError, 'Memory has been freed' if @freed
+        raise MemoryError, 'Write exceeds buffer size' if offset + data.bytesize > @size
+
+        @host_ptr[offset, data.bytesize] = data
+      end
+
+      # Read data from the pinned memory.
+      # @param count [Integer] Number of bytes to read
+      # @param offset [Integer] Offset in bytes
+      # @return [String]
+      def read(count, offset: 0)
+        raise MemoryError, 'Memory has been freed' if @freed
+        raise MemoryError, 'Read exceeds buffer size' if offset + count > @size
+
+        @host_ptr[offset, count]
+      end
+
+      # @return [Fiddle::Pointer]
+      def to_ptr
+        @host_ptr
+      end
+
+      class << self
+        # Create a finalizer for releasing pinned memory.
+        # @param ptr [Fiddle::Pointer]
+        # @return [Proc]
+        def release_finalizer(ptr)
+          ptr_addr = ptr.to_i
+          proc do
+            begin
+              RuntimeAPI.ensure_loaded!
+              RuntimeAPI.free_host(Fiddle::Pointer.new(ptr_addr))
+            rescue StandardError
+              # Silently ignore errors during finalization
+            end
+          end
+        end
+      end
+
+      private
+
+      # Allocate pinned host memory via RuntimeAPI.
+      # @return [Fiddle::Pointer]
+      def allocate_pinned_memory
+        RuntimeAPI.ensure_loaded!
+        RuntimeAPI.host_alloc(@size)
+      end
+    end
+  end
+end