ignis-collective 0.0.1

data/lib/nvruby/collective/net/rdma_transport.rb

@@ -0,0 +1,366 @@
+# frozen_string_literal: true
+
+require_relative "nd_bindings"
+require_relative "nd_adapter"
+require_relative "../transport/base"
+
+module Ignis
+  module Collective
+    module NetworkDirect
+      # RDMA Transport for multi-node GPU communication
+      #
+      # Uses Windows NetworkDirect for zero-copy, kernel-bypass transfers.
+      # Integrates with CUDA for GPU memory registration.
+      #
+      # Workflow:
+      # 1. Discover RDMA adapters
+      # 2. Create queue pairs and completion queues
+      # 3. Register GPU memory for RDMA
+      # 4. Connect to remote peers
+      # 5. Perform RDMA Read/Write operations
+      class RDMATransport < Transport::Base
+        # @return [Symbol] Transport type identifier
+        def self.transport_type
+          :rdma
+        end
+
+        # @return [Float] Estimated bandwidth (GB/s)
+        def estimated_bandwidth
+          100.0  # 100 Gbps Mellanox ConnectX-6
+        end
+
+        # @return [Float] Estimated latency (microseconds)
+        def estimated_latency
+          1.5  # RDMA latency typically < 2us
+        end
+
+        # Initialize the transport
+        # @param local_address [String] Local IP for RDMA bind
+        # @param local_port [Integer] Local port
+        def initialize(local_address: nil, local_port: nil, **opts)
+          super(**opts)
+          @local_address = local_address
+          @local_port = local_port || 0
+          @adapter = nil
+          @send_cq = nil
+          @recv_cq = nil
+          @qp = nil
+          @connector = nil
+          @memory_regions = {}  # gpu_ptr -> MemoryRegion
+        end
+
+        # Initialize RDMA resources
+        # @return [void]
+        def initialize!
+          return if @initialized
+
+          Bindings.ensure_loaded!
+          
+          unless Bindings.available?
+            raise RDMAError, "NetworkDirect not available: #{Bindings.load_error}"
+          end
+
+          # Start NetworkDirect
+          status = Bindings.NdStartup(2)  # Version 2
+          Bindings.check_status!(status, "NdStartup")
+
+          # Open first available adapter
+          @adapter = open_adapter
+
+          # Create completion queues
+          @send_cq = @adapter.create_completion_queue(depth: 256)
+          @recv_cq = @adapter.create_completion_queue(depth: 256)
+
+          # Create queue pair
+          @qp = @adapter.create_queue_pair(
+            send_cq: @send_cq,
+            recv_cq: @recv_cq,
+            send_depth: 64,
+            recv_depth: 64,
+            sge_count: 4
+          )
+
+          # Create connector
+          @connector = @adapter.create_connector
+
+          # Bind to local address
+          if @local_address
+            @connector.bind(address: @local_address, port: @local_port)
+          end
+
+          @initialized = true
+        end
+
+        # Connect to remote peer (client mode)
+        # @param remote_address [String] Remote IP address
+        # @param remote_port [Integer] Remote port
+        # @param private_data [String, nil] Connection private data
+        # @return [void]
+        def connect(remote_address:, remote_port:, private_data: nil)
+          ensure_initialized!
+
+          @connector.connect(
+            qp: @qp,
+            remote_address: remote_address,
+            remote_port: remote_port,
+            private_data: private_data
+          )
+        end
+
+        # Accept incoming connection (server mode)
+        # @param private_data [String, nil] Response private data
+        # @return [void]
+        def accept(private_data: nil)
+          ensure_initialized!
+
+          @connector.accept(qp: @qp, private_data: private_data)
+        end
+
+        # Register GPU memory for RDMA
+        # @param gpu_ptr [FFI::Pointer] GPU device pointer
+        # @param size [Integer] Size in bytes
+        # @return [Hash] Remote access info {:address, :token, :size}
+        def register_gpu_memory(gpu_ptr, size)
+          ensure_initialized!
+
+          # For GPU memory, we need to either:
+          # 1. Use cuMemExportToShareableHandle for mapping
+          # 2. Stage through pinned host memory
+
+          # Option 2 is more portable: allocate pinned host buffer
+          host_ptr = allocate_pinned_host(size)
+
+          # Register the host buffer with RDMA
+          mr = @adapter.register_memory(host_ptr, size)
+          @memory_regions[gpu_ptr.address] = {
+            memory_region: mr,
+            host_buffer: host_ptr,
+            size: size,
+            gpu_ptr: gpu_ptr
+          }
+
+          mr.remote_access_info
+        end
+
+        # Deregister GPU memory
+        # @param gpu_ptr [FFI::Pointer] GPU device pointer
+        # @return [void]
+        def deregister_gpu_memory(gpu_ptr)
+          info = @memory_regions.delete(gpu_ptr.address)
+          return unless info
+
+          info[:memory_region].deregister!
+          free_pinned_host(info[:host_buffer])
+        end
+
+        # RDMA Send (two-sided)
+        # @param buffer [FFI::Pointer] Source buffer (GPU)
+        # @param size [Integer] Size in bytes
+        # @return [void]
+        def send(buffer, size)
+          ensure_initialized!
+
+          info = @memory_regions[buffer.address]
+          raise RDMAError, "Buffer not registered" unless info
+
+          # Copy GPU -> Host (staging)
+          CUDA::RuntimeAPI.cudaMemcpy(
+            info[:host_buffer],
+            buffer,
+            size,
+            CUDA::RuntimeAPI::MEMCPY_DEVICE_TO_HOST
+          )
+
+          # Post RDMA send
+          @qp.post_send(
+            sge_list: [{
+              buffer: info[:host_buffer],
+              length: size,
+              token: info[:memory_region].token
+            }]
+          )
+
+          # Wait for completion
+          poll_completion(@send_cq)
+        end
+
+        # RDMA Receive (two-sided)
+        # @param buffer [FFI::Pointer] Destination buffer (GPU)
+        # @param size [Integer] Size in bytes
+        # @return [void]
+        def receive(buffer, size)
+          ensure_initialized!
+
+          info = @memory_regions[buffer.address]
+          raise RDMAError, "Buffer not registered" unless info
+
+          # Post RDMA receive
+          @qp.post_receive(
+            sge_list: [{
+              buffer: info[:host_buffer],
+              length: size,
+              token: info[:memory_region].token
+            }]
+          )
+
+          # Wait for completion
+          poll_completion(@recv_cq)
+
+          # Copy Host -> GPU
+          CUDA::RuntimeAPI.cudaMemcpy(
+            buffer,
+            info[:host_buffer],
+            size,
+            CUDA::RuntimeAPI::MEMCPY_HOST_TO_DEVICE
+          )
+        end
+
+        # RDMA Write (one-sided, zero-copy to remote)
+        # @param local_buffer [FFI::Pointer] Local GPU buffer
+        # @param remote_address [Integer] Remote buffer address
+        # @param remote_token [Integer] Remote memory token
+        # @param size [Integer] Size in bytes
+        # @return [void]
+        def rdma_write(local_buffer:, remote_address:, remote_token:, size:)
+          ensure_initialized!
+
+          info = @memory_regions[local_buffer.address]
+          raise RDMAError, "Buffer not registered" unless info
+
+          # Stage to host
+          CUDA::RuntimeAPI.cudaMemcpy(
+            info[:host_buffer],
+            local_buffer,
+            size,
+            CUDA::RuntimeAPI::MEMCPY_DEVICE_TO_HOST
+          )
+
+          # RDMA Write
+          @qp.rdma_write(
+            remote_address: remote_address,
+            remote_token: remote_token,
+            sge_list: [{
+              buffer: info[:host_buffer],
+              length: size,
+              token: info[:memory_region].token
+            }]
+          )
+
+          poll_completion(@send_cq)
+        end
+
+        # RDMA Read (one-sided, zero-copy from remote)
+        # @param local_buffer [FFI::Pointer] Local GPU buffer
+        # @param remote_address [Integer] Remote buffer address
+        # @param remote_token [Integer] Remote memory token
+        # @param size [Integer] Size in bytes
+        # @return [void]
+        def rdma_read(local_buffer:, remote_address:, remote_token:, size:)
+          ensure_initialized!
+
+          info = @memory_regions[local_buffer.address]
+          raise RDMAError, "Buffer not registered" unless info
+
+          # RDMA Read
+          @qp.rdma_read(
+            remote_address: remote_address,
+            remote_token: remote_token,
+            sge_list: [{
+              buffer: info[:host_buffer],
+              length: size,
+              token: info[:memory_region].token
+            }]
+          )
+
+          poll_completion(@send_cq)
+
+          # Stage to GPU
+          CUDA::RuntimeAPI.cudaMemcpy(
+            local_buffer,
+            info[:host_buffer],
+            size,
+            CUDA::RuntimeAPI::MEMCPY_HOST_TO_DEVICE
+          )
+        end
+
+        # Check if RDMA is available
+        # @return [Boolean] True if RDMA hardware present
+        def self.available?
+          Bindings.available?
+        end
+
+        # Cleanup
+        # @return [void]
+        def destroy!
+          # Deregister all memory
+          @memory_regions.each_key do |ptr_addr|
+            ptr = FFI::Pointer.new(:uint8, ptr_addr)
+            deregister_gpu_memory(ptr) rescue nil
+          end
+
+          @connector&.close!
+          @qp&.close!
+          @send_cq&.close!
+          @recv_cq&.close!
+          @adapter&.close!
+
+          Bindings.NdCleanup rescue nil
+
+          @initialized = false
+        end
+
+        private
+
+        def open_adapter
+          # Query available adapters
+          address_list_size = FFI::MemoryPointer.new(:uint32)
+          address_list_size.write_uint32(0)
+
+          # Get size first
+          Bindings.NdQueryAddressList(0, FFI::Pointer::NULL, address_list_size)
+
+          # Allocate and query
+          size = address_list_size.read_uint32
+          address_list = FFI::MemoryPointer.new(:uint8, size)
+          status = Bindings.NdQueryAddressList(0, address_list, address_list_size)
+          Bindings.check_status!(status, "NdQueryAddressList")
+
+          # Open first adapter
+          adapter_ptr = FFI::MemoryPointer.new(:pointer)
+          status = Bindings.NdOpenAdapter(address_list, size, adapter_ptr)
+          Bindings.check_status!(status, "NdOpenAdapter")
+
+          Adapter.new(adapter_ptr.read_pointer, "RDMA0")
+        end
+
+        def allocate_pinned_host(size)
+          ptr_ptr = FFI::MemoryPointer.new(:pointer)
+          status = CUDA::RuntimeAPI.cudaHostAlloc(ptr_ptr, size, 1)  # cudaHostAllocDefault
+          CUDA::RuntimeAPI.check_status!(status, "cudaHostAlloc for RDMA")
+          ptr_ptr.read_pointer
+        end
+
+        def free_pinned_host(ptr)
+          CUDA::RuntimeAPI.cudaFreeHost(ptr) rescue nil
+        end
+
+        def poll_completion(cq, timeout_ms: 5000)
+          start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          
+          loop do
+            results = cq.poll(max_results: 1)
+            return results.first if results.any?
+
+            elapsed = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
+            if elapsed > timeout_ms
+              raise RDMAError, "Completion poll timeout"
+            end
+
+            # Brief sleep to avoid busy-wait
+            sleep(0.0001)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/collective/nvarray_adapter.rb

@@ -0,0 +1,230 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    # Adapter for seamless NvArray integration with collective operations
+    # Handles dtype detection, shape validation, and buffer extraction
+    #
+    # @example Direct tensor usage
+    #   tensors = gpus.map { |gpu| NvArray.zeros([1000, 128], device: gpu) }
+    #   comm.all_reduce(tensors)  # NvArrayAdapter handles conversion
+    #
+    module NvArrayAdapter
+      # Supported dtypes and their byte sizes
+      DTYPE_SIZES = {
+        float32: 4,
+        float64: 8,
+        float16: 2,
+        bfloat16: 2,
+        int32: 4,
+        int64: 8,
+        int16: 2,
+        int8: 1,
+        uint8: 1,
+        uint32: 4,
+        uint64: 8
+      }.freeze
+
+      # CUDA dtype codes for kernel dispatch
+      DTYPE_CUDA_CODES = {
+        float32: 0,
+        float64: 1,
+        float16: 2,
+        bfloat16: 3,
+        int32: 4,
+        int64: 5,
+        int16: 6,
+        int8: 7,
+        uint8: 8,
+        uint32: 9,
+        uint64: 10
+      }.freeze
+
+      class << self
+        # Normalize input to array of NvArrays with validated properties
+        #
+        # @param inputs [Array<NvArray>, Array<FFI::Pointer>, NvArray] Input tensors
+        # @param expected_count [Integer, nil] Expected number of tensors (optional)
+        # @return [Array<NvArray>] Normalized tensor array
+        # @raise [ArgumentError] If inputs are invalid
+        def normalize(inputs, expected_count: nil)
+          tensors = wrap_array(inputs)
+
+          if expected_count && tensors.size != expected_count
+            raise ArgumentError,
+                  "Expected #{expected_count} tensors, got #{tensors.size}"
+          end
+
+          validate_tensors!(tensors)
+          tensors
+        end
+
+        # Extract device pointers from tensors for raw operations
+        #
+        # @param tensors [Array<NvArray>] Input tensors
+        # @return [Array<FFI::Pointer>] Device pointers
+        def extract_pointers(tensors)
+          tensors.map do |t|
+            if t.respond_to?(:data_ptr)
+              t.data_ptr
+            elsif t.respond_to?(:device_ptr)
+              t.device_ptr
+            elsif t.respond_to?(:pointer)
+              t.pointer
+            elsif t.is_a?(Fiddle::Pointer)
+              t
+            elsif defined?(FFI::Pointer) && t.is_a?(FFI::Pointer)
+              t
+            else
+              raise ArgumentError, "Cannot extract pointer from #{t.class}"
+            end
+          end
+        end
+
+        # Get common dtype from tensor array
+        #
+        # @param tensors [Array<NvArray>] Input tensors
+        # @return [Symbol] Common dtype
+        # @raise [ArgumentError] If dtypes don't match
+        def common_dtype(tensors)
+          return :float32 if tensors.empty?
+
+          dtypes = tensors.map { |t| extract_dtype(t) }.uniq
+          if dtypes.size > 1
+            raise ArgumentError,
+                  "All tensors must have same dtype, got: #{dtypes.join(', ')}"
+          end
+
+          dtypes.first
+        end
+
+        # Get element count from tensors (must match)
+        #
+        # @param tensors [Array<NvArray>] Input tensors
+        # @return [Integer] Element count
+        # @raise [ArgumentError] If element counts don't match
+        def common_element_count(tensors)
+          return 0 if tensors.empty?
+
+          counts = tensors.map { |t| extract_element_count(t) }.uniq
+          if counts.size > 1
+            raise ArgumentError,
+                  "All tensors must have same element count, got: #{counts.join(', ')}"
+          end
+
+          counts.first
+        end
+
+        # Get byte size for a dtype
+        #
+        # @param dtype [Symbol] Data type
+        # @return [Integer] Bytes per element
+        def dtype_size(dtype)
+          DTYPE_SIZES[dtype] || raise(ArgumentError, "Unknown dtype: #{dtype}")
+        end
+
+        # Get CUDA type code for kernel dispatch
+        #
+        # @param dtype [Symbol] Data type
+        # @return [Integer] CUDA type code
+        def dtype_cuda_code(dtype)
+          DTYPE_CUDA_CODES[dtype] || 0
+        end
+
+        # Calculate total byte size for tensors
+        #
+        # @param tensors [Array<NvArray>] Input tensors
+        # @return [Integer] Total bytes
+        def total_byte_size(tensors)
+          tensors.sum do |t|
+            extract_element_count(t) * dtype_size(extract_dtype(t))
+          end
+        end
+
+        # Validate shape compatibility for broadcast operations
+        #
+        # @param src_tensor [NvArray] Source tensor
+        # @param dst_tensors [Array<NvArray>] Destination tensors
+        # @return [Boolean] True if compatible
+        def broadcast_compatible?(src_tensor, dst_tensors)
+          src_shape = extract_shape(src_tensor)
+          dst_tensors.all? { |t| extract_shape(t) == src_shape }
+        end
+
+        # Create buffer info for collective ops
+        #
+        # @param tensors [Array<NvArray>] Input tensors
+        # @return [Hash] Buffer information
+        def buffer_info(tensors)
+          {
+            pointers: extract_pointers(tensors),
+            dtype: common_dtype(tensors),
+            dtype_code: dtype_cuda_code(common_dtype(tensors)),
+            element_count: common_element_count(tensors),
+            byte_size: common_element_count(tensors) * dtype_size(common_dtype(tensors)),
+            tensor_count: tensors.size
+          }
+        end
+
+        private
+
+        def wrap_array(inputs)
+          case inputs
+          when Array
+            inputs
+          when Ignis::NvArray
+            [inputs]
+          else
+            raise ArgumentError, "Expected Array or NvArray, got #{inputs.class}"
+          end
+        end
+
+        def validate_tensors!(tensors)
+          return if tensors.empty?
+
+          # All tensors must have extractable properties
+          tensors.each_with_index do |t, i|
+            unless t.respond_to?(:data_ptr) || t.respond_to?(:device_ptr) || t.respond_to?(:pointer) ||
+                   t.is_a?(Fiddle::Pointer) || (defined?(FFI::Pointer) && t.is_a?(FFI::Pointer))
+              raise ArgumentError, "Tensor #{i} has no data_ptr/device_ptr method"
+            end
+          end
+
+          # Check dtype consistency
+          common_dtype(tensors)
+
+          # Check element count consistency
+          common_element_count(tensors)
+        end
+
+        def extract_dtype(tensor)
+          if tensor.respond_to?(:dtype)
+            tensor.dtype
+          else
+            :float32  # Default assumption for raw pointers
+          end
+        end
+
+        def extract_element_count(tensor)
+          if tensor.respond_to?(:numel)
+            tensor.numel
+          elsif tensor.respond_to?(:shape)
+            tensor.shape.reduce(1, :*)
+          elsif tensor.respond_to?(:size)
+            tensor.size
+          else
+            raise ArgumentError, "Cannot determine element count for #{tensor.class}"
+          end
+        end
+
+        def extract_shape(tensor)
+          if tensor.respond_to?(:shape)
+            tensor.shape
+          else
+            nil
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/collective/p2p_bindings.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require "ffi"
+
+module Ignis
+  module Collective
+    # P2P-specific CUDA bindings for multi-GPU communication
+    # Extends RuntimeAPI with peer-to-peer and IPC functions
+    module P2PBindings
+      extend FFI::Library
+
+      # P2P attribute constants (cudaDeviceP2PAttr)
+      P2P_ATTR_PERFORMANCE_RANK = 1
+      P2P_ATTR_ACCESS_SUPPORTED = 2
+      P2P_ATTR_NATIVE_ATOMIC_SUPPORTED = 3
+      P2P_ATTR_CUDA_ARRAY_ACCESS_SUPPORTED = 4
+
+      # IPC memory flags
+      IPC_MEM_LAZY_ENABLE_PEER_ACCESS = 1
+
+      # CUDA IPC memory handle size (64 bytes)
+      IPC_HANDLE_SIZE = 64
+
+      # cudaIpcMemHandle_t structure
+      class CudaIpcMemHandle < FFI::Struct
+        layout :reserved, [:uint8, IPC_HANDLE_SIZE]
+      end
+
+      class << self
+        # Ensure P2P bindings are loaded
+        # @return [void]
+        def ensure_loaded!
+          return if @loaded
+
+          CUDA::LibraryLoader.ensure_cuda_runtime!
+
+          # Resolve cudart path per platform
+          dll_path = if defined?(Ignis::Platform)
+                       Ignis::Platform.cudart_path
+                     elsif RUBY_PLATFORM.match?(/mswin|mingw|cygwin/i)
+                       cuda_bin = Ignis.configuration.cuda_bin_path rescue nil
+                       cuda_bin ? Dir.glob(File.join(cuda_bin, 'cudart64_*.dll')).max : 'cudart64_130'
+                     else
+                       'libcudart.so.13'
+                     end
+
+          ffi_lib dll_path if dll_path
+
+          attach_p2p_functions!
+          @loaded = true
+        end
+
+        private
+
+        def attach_p2p_functions!
+          # Peer-to-peer access
+          attach_function :cudaDeviceCanAccessPeer,
+                          [:pointer, :int, :int],
+                          :int
+
+          attach_function :cudaDeviceEnablePeerAccess,
+                          [:int, :uint],
+                          :int
+
+          attach_function :cudaDeviceDisablePeerAccess,
+                          [:int],
+                          :int
+
+          attach_function :cudaDeviceGetP2PAttribute,
+                          [:pointer, :int, :int, :int],
+                          :int
+
+          # Peer-to-peer memory copy
+          attach_function :cudaMemcpyPeer,
+                          [:pointer, :int, :pointer, :int, :size_t],
+                          :int
+
+          attach_function :cudaMemcpyPeerAsync,
+                          [:pointer, :int, :pointer, :int, :size_t, :pointer],
+                          :int
+
+          # IPC memory handle functions
+          attach_function :cudaIpcGetMemHandle,
+                          [CudaIpcMemHandle.ptr, :pointer],
+                          :int
+
+          attach_function :cudaIpcOpenMemHandle,
+                          [:pointer, CudaIpcMemHandle.by_value, :uint],
+                          :int
+
+          attach_function :cudaIpcCloseMemHandle,
+                          [:pointer],
+                          :int
+
+          # IPC event handle functions
+          attach_function :cudaIpcGetEventHandle,
+                          [:pointer, :pointer],
+                          :int
+
+          attach_function :cudaIpcOpenEventHandle,
+                          [:pointer, :pointer],
+                          :int
+        end
+      end
+
+      # Check CUDA status and raise error if not success
+      # @param status [Integer] CUDA status code
+      # @param context [String] Error context
+      # @return [void]
+      # @raise [CudaRuntimeError] If status indicates an error
+      def self.check_status!(status, context = "P2P operation")
+        return if status.zero?
+
+        CUDA::RuntimeAPI.ensure_loaded!
+        error_name = CUDA::RuntimeAPI.cudaGetErrorName(status)
+        error_string = CUDA::RuntimeAPI.cudaGetErrorString(status)
+        raise CudaRuntimeError.new("#{context}: #{error_name} - #{error_string}", cuda_code: status)
+      end
+    end
+  end
+end