ignis 0.0.1

data/lib/nvruby/memory/device_memory_resource.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Memory
+    # @abstract Abstract base class for device memory resources
+    # Modeled after RMM's device_memory_resource interface
+    class DeviceMemoryResource
+      # Minimum alignment for all allocations (256 bytes per RMM standard)
+      ALIGNMENT = 256
+
+      # @return [Integer] Device index this resource manages
+      attr_reader :device_index
+
+      # @param device_index [Integer] GPU device index (default: current device)
+      # @return [void]
+      def initialize(device_index: nil)
+        @device_index = device_index || Ignis.configuration.default_device
+        @mutex = Mutex.new
+      end
+
+      # Allocate device memory with optional stream ordering
+      # @param bytes [Integer] Number of bytes to allocate
+      # @param stream [Ignis::CUDA::Stream, nil] Optional stream for async allocation
+      # @return [FFI::Pointer] Device pointer
+      # @raise [Ignis::MemoryError] If allocation fails
+      def allocate(bytes, stream: nil)
+        raise ArgumentError, "bytes must be positive, got #{bytes}" if bytes <= 0
+
+        aligned_bytes = align_up(bytes)
+
+        @mutex.synchronize do
+          ptr = do_allocate(aligned_bytes, stream)
+          Stats.record_allocation(aligned_bytes)
+          ptr
+        end
+      end
+
+      # Deallocate device memory with optional stream ordering
+      # @param ptr [FFI::Pointer] Device pointer to free
+      # @param bytes [Integer] Size of allocation
+      # @param stream [Ignis::CUDA::Stream, nil] Optional stream for async deallocation
+      # @return [void]
+      def deallocate(ptr, bytes, stream: nil)
+        return if ptr.null?
+
+        aligned_bytes = align_up(bytes)
+
+        @mutex.synchronize do
+          do_deallocate(ptr, aligned_bytes, stream)
+          Stats.record_deallocation(aligned_bytes)
+        end
+      end
+
+      # Check if this resource supports stream-ordered allocation
+      # @return [Boolean]
+      def supports_streams?
+        false
+      end
+
+      # Check if memory from this resource can be deallocated by another
+      # @param other [DeviceMemoryResource]
+      # @return [Boolean]
+      def is_equal?(other)
+        self.class == other.class && @device_index == other.device_index
+      end
+
+      # @return [String]
+      def to_s
+        "#{self.class.name}[device=#{@device_index}]"
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name}:0x#{object_id.to_s(16)} device=#{@device_index}>"
+      end
+
+      protected
+
+      # @abstract Subclasses must implement this
+      # @param bytes [Integer] Aligned bytes to allocate
+      # @param stream [Ignis::CUDA::Stream, nil]
+      # @return [FFI::Pointer]
+      def do_allocate(bytes, stream)
+        raise NotImplementedError, "#{self.class}#do_allocate not implemented"
+      end
+
+      # @abstract Subclasses must implement this
+      # @param ptr [FFI::Pointer]
+      # @param bytes [Integer]
+      # @param stream [Ignis::CUDA::Stream, nil]
+      # @return [void]
+      def do_deallocate(ptr, bytes, stream)
+        raise NotImplementedError, "#{self.class}#do_deallocate not implemented"
+      end
+
+      private
+
+      # Align size up to ALIGNMENT boundary
+      # @param bytes [Integer]
+      # @return [Integer]
+      def align_up(bytes)
+        ((bytes + ALIGNMENT - 1) / ALIGNMENT) * ALIGNMENT
+      end
+    end
+  end
+end

data/lib/nvruby/memory/pinned_host_memory_resource.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Memory
+    # @abstract Abstract base class for host memory resources
+    class HostMemoryResource
+      # Minimum alignment for all allocations
+      ALIGNMENT = 64
+
+      # @return [void]
+      def initialize
+        @mutex = Mutex.new
+      end
+
+      # Allocate host memory
+      # @param bytes [Integer] Number of bytes to allocate
+      # @return [FFI::Pointer] Host pointer
+      # @raise [Ignis::MemoryError] If allocation fails
+      def allocate(bytes)
+        raise ArgumentError, "bytes must be positive, got #{bytes}" if bytes <= 0
+
+        aligned_bytes = align_up(bytes)
+
+        @mutex.synchronize do
+          ptr = do_allocate(aligned_bytes)
+          Stats.record_allocation(aligned_bytes)
+          ptr
+        end
+      end
+
+      # Deallocate host memory
+      # @param ptr [FFI::Pointer] Host pointer to free
+      # @param bytes [Integer] Size of allocation
+      # @return [void]
+      def deallocate(ptr, bytes)
+        return if ptr.null?
+
+        aligned_bytes = align_up(bytes)
+
+        @mutex.synchronize do
+          do_deallocate(ptr, aligned_bytes)
+          Stats.record_deallocation(aligned_bytes)
+        end
+      end
+
+      protected
+
+      # @abstract
+      def do_allocate(bytes)
+        raise NotImplementedError
+      end
+
+      # @abstract
+      def do_deallocate(ptr, bytes)
+        raise NotImplementedError
+      end
+
+      private
+
+      def align_up(bytes)
+        ((bytes + ALIGNMENT - 1) / ALIGNMENT) * ALIGNMENT
+      end
+    end
+
+    # Pinned (page-locked) host memory resource
+    # Uses cudaHostAlloc for faster GPU transfers
+    class PinnedHostMemoryResource < HostMemoryResource
+      # cudaHostAllocDefault flag
+      HOST_ALLOC_DEFAULT = 0x00
+      # cudaHostAllocPortable flag
+      HOST_ALLOC_PORTABLE = 0x01
+      # cudaHostAllocMapped flag
+      HOST_ALLOC_MAPPED = 0x02
+      # cudaHostAllocWriteCombined flag
+      HOST_ALLOC_WRITE_COMBINED = 0x04
+
+      # @param flags [Integer] cudaHostAlloc flags (default: 0)
+      def initialize(flags: HOST_ALLOC_DEFAULT)
+        super()
+        @flags = flags
+        CUDA::RuntimeAPI.ensure_loaded!
+      end
+
+      protected
+
+      def do_allocate(bytes)
+        ptr_ptr = FFI::MemoryPointer.new(:pointer)
+
+        if @flags == HOST_ALLOC_DEFAULT
+          status = CUDA::RuntimeAPI.cudaMallocHost(ptr_ptr, bytes)
+          CUDA::RuntimeAPI.check_status!(status, "cudaMallocHost(#{bytes} bytes)")
+        else
+          status = cuda_host_alloc(ptr_ptr, bytes, @flags)
+          CUDA::RuntimeAPI.check_status!(status, "cudaHostAlloc(#{bytes} bytes, flags=#{@flags})")
+        end
+
+        ptr_ptr.read_pointer
+      end
+
+      def do_deallocate(ptr, _bytes)
+        status = CUDA::RuntimeAPI.cudaFreeHost(ptr)
+        CUDA::RuntimeAPI.check_status!(status, "cudaFreeHost")
+      end
+
+      private
+
+      def cuda_host_alloc(ptr_ptr, size, flags)
+        CUDA::RuntimeAPI.cudaHostAlloc(ptr_ptr, size, flags)
+      end
+    end
+  end
+end

data/lib/nvruby/memory/pool_memory_resource.rb

@@ -0,0 +1,242 @@
+# frozen_string_literal: true
+
+require_relative 'device_memory_resource'
+
+module Ignis
+  module Memory
+    # Coalescing best-fit pool memory resource
+    # Pre-allocates a large pool and sub-allocates from it
+    # Inspired by RMM's pool_memory_resource
+    class PoolMemoryResource < DeviceMemoryResource
+      # Default initial pool size: 256MB
+      DEFAULT_INITIAL_SIZE = 256 * 1024 * 1024
+
+      # Default maximum pool size: 90% of device memory
+      DEFAULT_MAX_RATIO = 0.9
+
+      # Block header for tracking allocations in the pool
+      Block = Struct.new(:ptr, :size, :free, keyword_init: true)
+
+      # @return [Integer] Current pool size
+      attr_reader :current_pool_size
+
+      # @return [Integer] Maximum pool size
+      attr_reader :maximum_pool_size
+
+      # @param upstream [DeviceMemoryResource] Upstream resource for pool allocations
+      # @param initial_pool_size [Integer] Initial pool size in bytes
+      # @param maximum_pool_size [Integer, nil] Maximum pool size (nil = 90% of device memory)
+      # @param device_index [Integer, nil] GPU device index
+      def initialize(
+        upstream: nil,
+        initial_pool_size: DEFAULT_INITIAL_SIZE,
+        maximum_pool_size: nil,
+        device_index: nil
+      )
+        super(device_index: device_index)
+
+        @upstream = upstream || CudaMemoryResource.new(device_index: @device_index)
+        @initial_pool_size = align_up(initial_pool_size)
+        @maximum_pool_size = maximum_pool_size || calculate_max_pool_size
+        @maximum_pool_size = align_up(@maximum_pool_size)
+
+        @current_pool_size = 0
+        @blocks = []
+        @free_list = []
+
+        allocate_initial_pool!
+      end
+
+      # @return [Boolean] false - synchronous pool allocation
+      def supports_streams?
+        false
+      end
+
+      # Release all pool memory back to upstream
+      # @return [void]
+      def destroy!
+        @mutex.synchronize do
+          @blocks.each do |block|
+            @upstream.deallocate(block.ptr, block.size, stream: nil)
+          end
+          @blocks.clear
+          @free_list.clear
+          @current_pool_size = 0
+        end
+      end
+
+      # @return [Hash] Pool statistics
+      def pool_stats
+        @mutex.synchronize do
+          total_free = @free_list.sum(&:size)
+          {
+            current_pool_size: @current_pool_size,
+            maximum_pool_size: @maximum_pool_size,
+            free_bytes: total_free,
+            used_bytes: @current_pool_size - total_free,
+            block_count: @blocks.size,
+            free_block_count: @free_list.size
+          }
+        end
+      end
+
+      protected
+
+      # @param bytes [Integer]
+      # @param stream [Ignis::CUDA::Stream, nil] Ignored for pool allocation
+      # @return [FFI::Pointer]
+      def do_allocate(bytes, _stream)
+        block = find_free_block(bytes)
+
+        if block
+          split_block_if_needed!(block, bytes)
+          block.ptr
+        else
+          expand_pool!(bytes)
+          block = find_free_block(bytes)
+          raise MemoryError, "Pool allocation failed for #{bytes} bytes" unless block
+
+          split_block_if_needed!(block, bytes)
+          block.ptr
+        end
+      end
+
+      # @param ptr [FFI::Pointer]
+      # @param bytes [Integer]
+      # @param stream [Ignis::CUDA::Stream, nil] Ignored for pool deallocation
+      # @return [void]
+      def do_deallocate(ptr, bytes, _stream)
+        block = @blocks.find { |b| b.ptr.address == ptr.address }
+        return unless block
+
+        block.free = true
+        @free_list << block
+        coalesce_free_blocks!
+      end
+
+      private
+
+      def allocate_initial_pool!
+        return if @initial_pool_size <= 0
+
+        ptr = @upstream.allocate(@initial_pool_size, stream: nil)
+        block = Block.new(ptr: ptr, size: @initial_pool_size, free: true)
+        @blocks << block
+        @free_list << block
+        @current_pool_size = @initial_pool_size
+
+        Ignis.logger.debug { "Pool initialized with #{format_bytes(@initial_pool_size)}" }
+      end
+
+      def find_free_block(bytes)
+        @free_list
+          .select { |b| b.free && b.size >= bytes }
+          .min_by(&:size)
+      end
+
+      def split_block_if_needed!(block, bytes)
+        block.free = false
+        @free_list.delete(block)
+
+        remainder = block.size - bytes
+        return if remainder < ALIGNMENT * 2
+
+        new_ptr = FFI::Pointer.new(block.ptr.address + bytes)
+        new_block = Block.new(ptr: new_ptr, size: remainder, free: true)
+        block.size = bytes
+
+        @blocks << new_block
+        @free_list << new_block
+      end
+
+      def expand_pool!(min_bytes)
+        if @current_pool_size >= @maximum_pool_size
+          raise MemoryError, "Pool exhausted: current=#{format_bytes(@current_pool_size)}, " \
+                            "max=#{format_bytes(@maximum_pool_size)}, requested=#{format_bytes(min_bytes)}"
+        end
+
+        remaining = @maximum_pool_size - @current_pool_size
+        grow_size = if @maximum_pool_size > 0
+                      [remaining / 2, min_bytes, ALIGNMENT].max
+                    else
+                      [@current_pool_size, min_bytes, @initial_pool_size].max
+                    end
+        grow_size = align_up([grow_size, remaining].min)
+
+        Ignis.logger.debug { "Expanding pool by #{format_bytes(grow_size)}" }
+
+        ptr = @upstream.allocate(grow_size, stream: nil)
+        block = Block.new(ptr: ptr, size: grow_size, free: true)
+        @blocks << block
+        @free_list << block
+        @current_pool_size += grow_size
+      end
+
+      def coalesce_free_blocks!
+        return if @free_list.size < 2
+
+        sorted = @free_list.sort_by { |b| b.ptr.address }
+        merged = []
+        current = sorted.first
+
+        sorted.drop(1).each do |block|
+          if current.ptr.address + current.size == block.ptr.address
+            current.size += block.size
+            @blocks.delete(block)
+          else
+            merged << current
+            current = block
+          end
+        end
+        merged << current
+
+        @free_list.replace(merged)
+      end
+
+      def calculate_max_pool_size
+        CUDA::RuntimeAPI.ensure_loaded!
+        free_ptr = FFI::MemoryPointer.new(:size_t)
+        total_ptr = FFI::MemoryPointer.new(:size_t)
+
+        ensure_device do
+          status = CUDA::RuntimeAPI.cudaMemGetInfo(free_ptr, total_ptr)
+          CUDA::RuntimeAPI.check_status!(status, "cudaMemGetInfo")
+        end
+
+        total = total_ptr.read(:size_t)
+        (total * DEFAULT_MAX_RATIO).to_i
+      end
+
+      def ensure_device
+        current_ptr = FFI::MemoryPointer.new(:int)
+        CUDA::RuntimeAPI.cudaGetDevice(current_ptr)
+        original = current_ptr.read_int
+
+        if original != @device_index
+          status = CUDA::RuntimeAPI.cudaSetDevice(@device_index)
+          CUDA::RuntimeAPI.check_status!(status, "cudaSetDevice(#{@device_index})")
+        end
+
+        result = yield
+
+        if original != @device_index
+          CUDA::RuntimeAPI.cudaSetDevice(original)
+        end
+
+        result
+      end
+
+      def format_bytes(bytes)
+        if bytes >= 1024 * 1024 * 1024
+          format("%.2f GB", bytes.to_f / (1024 * 1024 * 1024))
+        elsif bytes >= 1024 * 1024
+          format("%.2f MB", bytes.to_f / (1024 * 1024))
+        elsif bytes >= 1024
+          format("%.2f KB", bytes.to_f / 1024)
+        else
+          "#{bytes} B"
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/memory/stats.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Memory
+    # Thread-safe allocation statistics tracking
+    # Provides real-time visibility into memory usage patterns
+    class Stats
+      class << self
+        # @return [Integer] Total bytes currently allocated
+        def total_allocated_bytes
+          @total_allocated_bytes ||= 0
+        end
+
+        # @return [Integer] Peak bytes allocated (high-water mark)
+        def peak_allocated_bytes
+          @peak_allocated_bytes ||= 0
+        end
+
+        # @return [Integer] Total allocation count since startup
+        def allocation_count
+          @allocation_count ||= 0
+        end
+
+        # @return [Integer] Total deallocation count since startup
+        def deallocation_count
+          @deallocation_count ||= 0
+        end
+
+        # Record an allocation
+        # @param bytes [Integer]
+        # @return [void]
+        def record_allocation(bytes)
+          mutex.synchronize do
+            @total_allocated_bytes = (total_allocated_bytes + bytes)
+            @peak_allocated_bytes = [@peak_allocated_bytes || 0, @total_allocated_bytes].max
+            @allocation_count = (allocation_count + 1)
+          end
+        end
+
+        # Record a deallocation
+        # @param bytes [Integer]
+        # @return [void]
+        def record_deallocation(bytes)
+          mutex.synchronize do
+            @total_allocated_bytes = [total_allocated_bytes - bytes, 0].max
+            @deallocation_count = (deallocation_count + 1)
+          end
+        end
+
+        # Reset all statistics
+        # @return [void]
+        def reset!
+          mutex.synchronize do
+            @total_allocated_bytes = 0
+            @peak_allocated_bytes = 0
+            @allocation_count = 0
+            @deallocation_count = 0
+          end
+        end
+
+        # @return [Hash] Statistics snapshot
+        def snapshot
+          mutex.synchronize do
+            {
+              total_allocated_bytes: total_allocated_bytes,
+              peak_allocated_bytes: peak_allocated_bytes,
+              allocation_count: allocation_count,
+              deallocation_count: deallocation_count,
+              active_allocations: allocation_count - deallocation_count
+            }
+          end
+        end
+
+        # @return [String] Human-readable statistics
+        def to_s
+          s = snapshot
+          format(
+            "Memory Stats: %s allocated (peak: %s), %d allocs, %d frees, %d active",
+            format_bytes(s[:total_allocated_bytes]),
+            format_bytes(s[:peak_allocated_bytes]),
+            s[:allocation_count],
+            s[:deallocation_count],
+            s[:active_allocations]
+          )
+        end
+
+        private
+
+        def mutex
+          @mutex ||= Mutex.new
+        end
+
+        def format_bytes(bytes)
+          if bytes >= 1024 * 1024 * 1024
+            format("%.2f GB", bytes.to_f / (1024 * 1024 * 1024))
+          elsif bytes >= 1024 * 1024
+            format("%.2f MB", bytes.to_f / (1024 * 1024))
+          elsif bytes >= 1024
+            format("%.2f KB", bytes.to_f / 1024)
+          else
+            "#{bytes} B"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/memory.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+require_relative 'memory/stats'
+require_relative 'memory/device_memory_resource'
+require_relative 'memory/cuda_memory_resource'
+require_relative 'memory/cuda_async_memory_resource'
+require_relative 'memory/pool_memory_resource'
+require_relative 'memory/pinned_host_memory_resource'
+
+module Ignis
+  # Memory management module providing RMM-inspired GPU memory pools
+  #
+  # This module provides efficient memory allocation for GPU workloads through:
+  # - DeviceMemoryResource: Base class for all device memory allocators
+  # - CudaMemoryResource: Simple cudaMalloc/cudaFree (baseline)
+  # - CudaAsyncMemoryResource: Stream-ordered cudaMallocAsync/cudaFreeAsync
+  # - PoolMemoryResource: Coalescing best-fit pool for high-frequency allocations
+  # - PinnedHostMemoryResource: Page-locked host memory for fast transfers
+  #
+  # @example Using a memory pool
+  #   # Set global pool for all allocations
+  #   pool = Ignis::Memory::PoolMemoryResource.new(initial_pool_size: 512.megabytes)
+  #   Ignis::Memory.set_current_device_resource(pool)
+  #
+  #   # All NvArray allocations now use the pool
+  #   arr = Ignis::NvArray.new(shape: [1024, 1024], dtype: :float32)
+  #
+  # @example Stream-ordered allocation
+  #   async_mr = Ignis::Memory::CudaAsyncMemoryResource.new
+  #   ptr = async_mr.allocate(1024 * 1024, stream: my_stream)
+  #   # ...use ptr...
+  #   async_mr.deallocate(ptr, 1024 * 1024, stream: my_stream)
+  #
+  module Memory
+    class << self
+      # Get the current device memory resource
+      # Creates a default PoolMemoryResource if none is set
+      # @return [DeviceMemoryResource]
+      def get_current_device_resource
+        device_id = current_device_index
+        resource = per_device_resources[device_id]
+        return resource if resource
+
+        if Ignis.configuration.use_memory_pool
+          resource = PoolMemoryResource.new(device_index: device_id)
+        else
+          resource = CudaMemoryResource.new(device_index: device_id)
+        end
+
+        per_device_resources[device_id] = resource
+        resource
+      end
+
+      # Set the current device memory resource
+      # @param resource [DeviceMemoryResource]
+      # @return [DeviceMemoryResource] The previous resource
+      def set_current_device_resource(resource)
+        device_id = current_device_index
+        old = per_device_resources[device_id]
+        per_device_resources[device_id] = resource
+        old
+      end
+
+      # Get memory resource for a specific device
+      # @param device_id [Integer]
+      # @return [DeviceMemoryResource, nil]
+      def get_per_device_resource(device_id)
+        per_device_resources[device_id]
+      end
+
+      # Set memory resource for a specific device
+      # @param device_id [Integer]
+      # @param resource [DeviceMemoryResource]
+      # @return [DeviceMemoryResource] The previous resource
+      def set_per_device_resource(device_id, resource)
+        old = per_device_resources[device_id]
+        per_device_resources[device_id] = resource
+        old
+      end
+
+      # Get memory statistics
+      # @return [Hash]
+      def stats
+        Stats.snapshot
+      end
+
+      # Reset memory statistics
+      # @return [void]
+      def reset_stats!
+        Stats.reset!
+      end
+
+      # Allocate device memory using current resource
+      # @param bytes [Integer]
+      # @param stream [Ignis::CUDA::Stream, nil]
+      # @return [FFI::Pointer]
+      def allocate(bytes, stream: nil)
+        get_current_device_resource.allocate(bytes, stream: stream)
+      end
+
+      # Deallocate device memory using current resource
+      # @param ptr [FFI::Pointer]
+      # @param bytes [Integer]
+      # @param stream [Ignis::CUDA::Stream, nil]
+      # @return [void]
+      def deallocate(ptr, bytes, stream: nil)
+        get_current_device_resource.deallocate(ptr, bytes, stream: stream)
+      end
+
+      private
+
+      def per_device_resources
+        @per_device_resources ||= {}
+      end
+
+      def current_device_index
+        CUDA::RuntimeAPI.ensure_loaded!
+        ptr = FFI::MemoryPointer.new(:int)
+        CUDA::RuntimeAPI.cudaGetDevice(ptr)
+        ptr.read_int
+      end
+    end
+  end
+end

data/lib/nvruby/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Ignis
+  VERSION = "0.1.0"
+end