ignis-collective 0.0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 7343f5cf832fab47fa86cc79eab930cafb142b6f0967687d65d8b5896c8ad488
+  data.tar.gz: 87f2e47d253db19ede66f71d9fcfc24bc59fc1b7813b67230d6fe1bd1be87b67
+SHA512:
+  metadata.gz: 791c617270822aad2c14bbca288da77baf1b3202c5ed46d0695936ec8953582d1c6e0677e513c58f21ebde11ef07eacb795f480cc126a5c337c59e52943f38fd
+  data.tar.gz: 78a9a949ca36e3b880c7df0ebe7ccdc3408ed082546957703936ef80a810f27d5c98e74ee5b047a76770bc31105bb44326f76519d57d395e8aadb2ea07563602

data/README.md

@@ -0,0 +1,7 @@
+# ignis-collective  (experimental)
+
+Multi-GPU collective communication for Ruby — NCCL-style ring/tree all-reduce with P2P / IPC (VMM) / host-staged / TCP transports, on the [`ignis`](https://rubygems.org/gems/ignis) foundation.
+
+> **Experimental.** The transports require multiple GPUs / nodes and cannot be exercised on a single GPU, so this gem is shipped separately from the verified single-GPU stack. APIs may change.
+
+MIT.

data/lib/ignis-collective.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+# ignis-collective — multi-GPU collective communication (NvCCL: ring/tree all-reduce,
+# P2P/IPC/host-staged/TCP transports) on the Ignis foundation. EXPERIMENTAL: the
+# transports cannot be exercised on a single GPU, so this gem is provided for
+# multi-GPU/multi-node setups and is not part of the verified single-GPU path.
+
+require "ignis"
+require_relative "nvruby/collective"

data/lib/nvruby/collective/algorithms/double_binary_tree.rb

@@ -0,0 +1,364 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    module Algorithms
+      # Double Binary Tree Algorithm for heterogeneous GPU counts
+      #
+      # When GPU count is not a power of 2, a single binary tree has imbalanced load.
+      # The double binary tree uses two overlapping trees to balance communication.
+      #
+      # Key insight: GPU i participates in both trees but at different levels.
+      # - Tree 1: Standard binary tree rooted at 0
+      # - Tree 2: Binary tree rooted at N/2 (or closest power of 2)
+      #
+      # This achieves near-optimal latency even for non-power-of-2 GPU counts.
+      class DoubleBinaryTree
+        # Node in the double tree structure
+        TreeNode = Struct.new(:rank, :gpu_id, :parent, :left_child, :right_child, :tree_id, keyword_init: true)
+
+        # @return [Array<Integer>] GPU IDs
+        attr_reader :gpu_ids
+
+        # @return [Integer] Number of GPUs
+        attr_reader :n_gpus
+
+        # @return [TransportSelector] Transport selector
+        attr_reader :transport_selector
+
+        # @param gpu_ids [Array<Integer>] GPU device IDs
+        # @param transport_selector [TransportSelector] Transport selector
+        def initialize(gpu_ids:, transport_selector:)
+          @gpu_ids = gpu_ids.dup.freeze
+          @n_gpus = gpu_ids.size
+          @transport_selector = transport_selector
+          @tree1 = nil
+          @tree2 = nil
+        end
+
+        # Build the double binary tree structure
+        # @return [void]
+        def build!
+          @tree1 = build_binary_tree(root_rank: 0, tree_id: 1)
+          
+          # For tree 2, use a different root to balance load
+          # Root is at opposite end of the array
+          tree2_root = @n_gpus / 2
+          @tree2 = build_binary_tree(root_rank: tree2_root, tree_id: 2)
+        end
+
+        # Broadcast using double tree (for fault tolerance / load balancing)
+        #
+        # @param buffer [FFI::Pointer] Source buffer on root
+        # @param buffers [Array<FFI::Pointer>] Destination buffers on all GPUs
+        # @param size [Integer] Buffer size in bytes
+        # @param root [Integer] Root rank
+        # @param streams [Array<FFI::Pointer>] CUDA streams
+        # @return [void]
+        def broadcast(buffer:, buffers:, size:, root:, streams:)
+          return if @n_gpus == 1
+          
+          build! if @tree1.nil?
+
+          # For non-power-of-2, use split strategy
+          if power_of_2?(@n_gpus)
+            # Standard single tree broadcast
+            broadcast_tree(@tree1, buffers, size, root, streams)
+          else
+            # Split data: half via tree1, half via tree2
+            half_size = size / 2
+            
+            # Tree 1 handles first half
+            broadcast_tree_partial(@tree1, buffers, 0, half_size, root, streams)
+            
+            # Tree 2 handles second half
+            broadcast_tree_partial(@tree2, buffers, half_size, half_size, root, streams)
+          end
+        end
+
+        # Reduce using double tree
+        #
+        # @param buffers [Array<FFI::Pointer>] Source buffers
+        # @param sizes [Array<Integer>] Buffer sizes
+        # @param dtype [Symbol] Data type
+        # @param op [Symbol] Reduction operation
+        # @param root [Integer] Root rank
+        # @param streams [Array<FFI::Pointer>] CUDA streams
+        # @return [void]
+        def reduce(buffers:, sizes:, dtype:, op:, root:, streams:)
+          return if @n_gpus == 1
+          
+          build! if @tree1.nil?
+
+          if power_of_2?(@n_gpus)
+            reduce_tree(@tree1, buffers, sizes[0], dtype, op, root, streams)
+          else
+            # Split reduction via both trees, combine at root
+            half_size = sizes[0] / 2
+            elem_size = dtype_elem_size(dtype)
+            half_count = half_size / elem_size
+
+            # Allocate temp buffer at root for tree2 partial result
+            temp_buffer = allocate_buffer(@gpu_ids[root], half_size)
+
+            begin
+              # Tree 1: reduce first half to root
+              reduce_tree_partial(@tree1, buffers, 0, half_size, dtype, op, root, streams)
+              
+              # Tree 2: reduce second half to tree2's root, then send to actual root
+              tree2_root = @n_gpus / 2
+              reduce_tree_partial(@tree2, buffers, half_size, half_size, dtype, op, tree2_root, streams)
+              
+              # If tree2 root != actual root, transfer
+              if tree2_root != root
+                transport = @transport_selector.select_transport(@gpu_ids[tree2_root], @gpu_ids[root])
+                if transport.is_a?(Transport::P2PTransport)
+                  src = ptr_offset(buffers[tree2_root], half_size)
+                  dst = ptr_offset(buffers[root], half_size)
+                  transport.copy_async(dst, src, half_size, get_stream_ptr(streams[root]))
+                end
+              end
+              
+              synchronize_all_streams!(streams)
+            ensure
+              free_buffer(temp_buffer, @gpu_ids[root])
+            end
+          end
+        end
+
+        private
+
+        def build_binary_tree(root_rank:, tree_id:)
+          nodes = @n_gpus.times.map do |rank|
+            TreeNode.new(
+              rank: rank,
+              gpu_id: @gpu_ids[rank],
+              parent: nil,
+              left_child: nil,
+              right_child: nil,
+              tree_id: tree_id
+            )
+          end
+
+          # Build parent-child relationships
+          # Reorder so root is at position 0 in tree
+          reordered = reorder_for_root(root_rank)
+          
+          reordered.each_with_index do |orig_rank, tree_pos|
+            node = nodes[orig_rank]
+            
+            left_tree_pos = 2 * tree_pos + 1
+            right_tree_pos = 2 * tree_pos + 2
+            
+            if left_tree_pos < @n_gpus
+              left_orig_rank = reordered[left_tree_pos]
+              node.left_child = left_orig_rank
+              nodes[left_orig_rank].parent = orig_rank
+            end
+            
+            if right_tree_pos < @n_gpus
+              right_orig_rank = reordered[right_tree_pos]
+              node.right_child = right_orig_rank
+              nodes[right_orig_rank].parent = orig_rank
+            end
+          end
+
+          nodes
+        end
+
+        def reorder_for_root(root_rank)
+          # Circular shift so root is at index 0
+          order = (0...@n_gpus).to_a
+          order.rotate(root_rank)
+        end
+
+        def broadcast_tree(tree, buffers, size, root, streams)
+          depth = tree_depth
+          
+          depth.times do |d|
+            tree.each do |node|
+              next unless node_depth(node, tree) == d
+              
+              [node.left_child, node.right_child].compact.each do |child_rank|
+                src_gpu = node.gpu_id
+                dst_gpu = @gpu_ids[child_rank]
+                
+                transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+                stream_ptr = get_stream_ptr(streams[node.rank])
+                
+                if transport.is_a?(Transport::P2PTransport)
+                  transport.copy_async(buffers[child_rank], buffers[node.rank], size, stream_ptr)
+                end
+              end
+            end
+            
+            synchronize_all_streams!(streams)
+          end
+        end
+
+        def broadcast_tree_partial(tree, buffers, offset, size, root, streams)
+          depth = tree_depth
+          
+          depth.times do |d|
+            tree.each do |node|
+              next unless node_depth(node, tree) == d
+              
+              [node.left_child, node.right_child].compact.each do |child_rank|
+                src_gpu = node.gpu_id
+                dst_gpu = @gpu_ids[child_rank]
+                
+                transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+                stream_ptr = get_stream_ptr(streams[node.rank])
+                
+                if transport.is_a?(Transport::P2PTransport)
+                  src = ptr_offset(buffers[node.rank], offset)
+                  dst = ptr_offset(buffers[child_rank], offset)
+                  transport.copy_async(dst, src, size, stream_ptr)
+                end
+              end
+            end
+            
+            synchronize_all_streams!(streams)
+          end
+        end
+
+        def reduce_tree(tree, buffers, size, dtype, op, root, streams)
+          depth = tree_depth
+          elem_size = dtype_elem_size(dtype)
+          count = size / elem_size
+          
+          recv_buffers = allocate_recv_buffers(size)
+          
+          begin
+            (depth - 1).downto(0) do |d|
+              tree.each do |node|
+                next unless node_depth(node, tree) == d
+                next unless node.left_child || node.right_child
+                
+                CUDA::RuntimeAPI.cudaSetDevice(node.gpu_id)
+                stream_ptr = get_stream_ptr(streams[node.rank])
+                
+                [node.left_child, node.right_child].compact.each do |child_rank|
+                  src_gpu = @gpu_ids[child_rank]
+                  
+                  transport = @transport_selector.select_transport(src_gpu, node.gpu_id)
+                  
+                  if transport.is_a?(Transport::P2PTransport)
+                    transport.copy_async(recv_buffers[node.rank], buffers[child_rank], size, stream_ptr)
+                  end
+                  
+                  synchronize_stream!(streams[node.rank])
+                  
+                  ReductionOps.execute(op, buffers[node.rank], recv_buffers[node.rank], 
+                                       buffers[node.rank], count, dtype, stream_ptr)
+                end
+              end
+              
+              synchronize_all_streams!(streams)
+            end
+          ensure
+            free_recv_buffers(recv_buffers)
+          end
+        end
+
+        def reduce_tree_partial(tree, buffers, offset, size, dtype, op, root, streams)
+          # Similar to reduce_tree but only for a portion of the buffer
+          # Implementation mirrors reduce_tree with offset applied
+          reduce_tree(
+            tree,
+            buffers.map { |b| ptr_offset(b, offset) },
+            size,
+            dtype,
+            op,
+            root,
+            streams
+          )
+        end
+
+        def tree_depth
+          return 0 if @n_gpus <= 1
+          (Math.log2(@n_gpus).ceil)
+        end
+
+        def node_depth(node, tree)
+          depth = 0
+          current = node
+          while current.parent
+            current = tree[current.parent]
+            depth += 1
+          end
+          depth
+        end
+
+        def power_of_2?(n)
+          n > 0 && (n & (n - 1)).zero?
+        end
+
+        def dtype_elem_size(dtype)
+          case dtype
+          when :float32, :int32 then 4
+          when :float64, :int64 then 8
+          when :float16, :bfloat16 then 2
+          else 4
+          end
+        end
+
+        def get_stream_ptr(stream)
+          case stream
+          when FFI::Pointer then stream
+          when CUDA::Stream then stream.ptr
+          else FFI::Pointer::NULL
+          end
+        end
+
+        def ptr_offset(ptr, offset)
+          FFI::Pointer.new(:uint8, ptr.address + offset)
+        end
+
+        def allocate_recv_buffers(size)
+          @gpu_ids.map do |gpu_id|
+            allocate_buffer(gpu_id, size)
+          end
+        end
+
+        def allocate_buffer(gpu_id, size)
+          CUDA::RuntimeAPI.ensure_loaded!
+          CUDA::RuntimeAPI.cudaSetDevice(gpu_id)
+          ptr_ptr = FFI::MemoryPointer.new(:pointer)
+          CUDA::RuntimeAPI.cudaMalloc(ptr_ptr, size)
+          ptr_ptr.read_pointer
+        end
+
+        def free_buffer(buf, gpu_id)
+          return unless buf && !buf.null?
+          CUDA::RuntimeAPI.cudaSetDevice(gpu_id)
+          CUDA::RuntimeAPI.cudaFree(buf)
+        rescue StandardError
+          # Ignore cleanup errors
+        end
+
+        def free_recv_buffers(buffers)
+          buffers.each_with_index do |buf, i|
+            free_buffer(buf, @gpu_ids[i])
+          end
+        end
+
+        def synchronize_stream!(stream)
+          stream_ptr = get_stream_ptr(stream)
+          if stream_ptr.null?
+            CUDA::RuntimeAPI.cudaDeviceSynchronize
+          else
+            CUDA::RuntimeAPI.cudaStreamSynchronize(stream_ptr)
+          end
+        end
+
+        def synchronize_all_streams!(streams)
+          streams.each_with_index do |stream, i|
+            CUDA::RuntimeAPI.cudaSetDevice(@gpu_ids[i])
+            synchronize_stream!(stream)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/collective/algorithms/pipeliner.rb

@@ -0,0 +1,222 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    module Algorithms
+      # Message Pipelining for Overlapping Communication and Computation
+      #
+      # Splits large messages into chunks that can be processed concurrently:
+      # - While GPU processes chunk N, transfer chunk N+1
+      # - Hides transfer latency behind compute
+      #
+      # Optimal chunk size balances:
+      # - Too small: overhead dominates
+      # - Too large: no overlap opportunity
+      class Pipeliner
+        # Pipeline stage
+        PipelineStage = Struct.new(:chunk_id, :offset, :size, :state, keyword_init: true)
+        
+        # Pipeline states
+        STATE_PENDING = :pending
+        STATE_TRANSFERRING = :transferring
+        STATE_COMPUTING = :computing
+        STATE_COMPLETE = :complete
+
+        # Default chunk size: 512KB (good for PCIe 4.0)
+        DEFAULT_CHUNK_SIZE = 512 * 1024
+
+        # @return [Integer] Chunk size in bytes
+        attr_reader :chunk_size
+
+        # @return [Integer] Number of pipeline stages (double-buffering = 2)
+        attr_reader :num_stages
+
+        # @param chunk_size [Integer] Size of each pipeline chunk
+        # @param num_stages [Integer] Number of concurrent stages (default 2)
+        def initialize(chunk_size: DEFAULT_CHUNK_SIZE, num_stages: 2)
+          @chunk_size = chunk_size
+          @num_stages = num_stages
+          @streams = {}
+        end
+
+        # Calculate optimal chunk size based on bandwidth and compute speed
+        #
+        # @param bandwidth_gbps [Float] Transfer bandwidth in GB/s
+        # @param compute_gflops [Float] Compute throughput in GFLOPS
+        # @param flops_per_element [Float] FLOPs per element in reduction
+        # @param element_size [Integer] Bytes per element
+        # @return [Integer] Optimal chunk size in bytes
+        def self.optimal_chunk_size(bandwidth_gbps:, compute_gflops:, flops_per_element: 1.0, element_size: 4)
+          # Balance: transfer_time = compute_time
+          # chunk_size / bandwidth = (chunk_size / element_size) * flops_per_element / compute_throughput
+          
+          # Convert GB/s to bytes/s
+          bandwidth_bps = bandwidth_gbps * 1e9
+          # Convert GFLOPS to FLOPS
+          compute_flops = compute_gflops * 1e9
+
+          # chunk_size / bandwidth = (chunk_size * flops_per_element) / (element_size * compute_flops)
+          # Solving: chunk_size = bandwidth * element_size * compute_flops / (compute_flops + bandwidth * flops_per_element)
+          
+          # Simplified: aim for transfer time ~ compute time
+          # chunk ~= bandwidth * target_latency_seconds
+          target_latency = 0.001  # 1ms pipeline stage
+          optimal = (bandwidth_bps * target_latency).to_i
+
+          # Clamp to reasonable range
+          [[optimal, 64 * 1024].max, 4 * 1024 * 1024].min
+        end
+
+        # Execute a pipelined reduction operation
+        #
+        # @param buffers [Array<FFI::Pointer>] Device buffers
+        # @param total_size [Integer] Total buffer size
+        # @param dtype [Symbol] Data type
+        # @param op [Symbol] Reduction operation
+        # @param transfer_fn [Proc] Transfer function (chunk_offset, chunk_size, stream) -> void
+        # @param reduce_fn [Proc] Reduce function (chunk_offset, chunk_size, stream) -> void
+        # @param streams [Array<CUDA::Stream>] CUDA streams (one per stage)
+        # @return [void]
+        def execute_pipelined(buffers:, total_size:, dtype:, op:, transfer_fn:, reduce_fn:, streams: nil)
+          # Calculate chunks
+          n_chunks = (total_size + @chunk_size - 1) / @chunk_size
+          
+          # Create streams if not provided
+          streams ||= create_streams(@num_stages, buffers[0])
+          
+          # Initialize pipeline stages
+          stages = n_chunks.times.map do |i|
+            offset = i * @chunk_size
+            size = [total_size - offset, @chunk_size].min
+            PipelineStage.new(chunk_id: i, offset: offset, size: size, state: STATE_PENDING)
+          end
+
+          # Pipeline execution
+          active_stages = []
+          pending_stages = stages.dup
+          
+          while pending_stages.any? || active_stages.any?
+            # Start new transfers up to num_stages
+            while active_stages.size < @num_stages && pending_stages.any?
+              stage = pending_stages.shift
+              stream = streams[stage.chunk_id % @num_stages]
+              
+              # Start transfer
+              transfer_fn.call(stage.offset, stage.size, stream)
+              stage.state = STATE_TRANSFERRING
+              active_stages << stage
+            end
+
+            # Check for completed transfers, start compute
+            active_stages.each do |stage|
+              if stage.state == STATE_TRANSFERRING
+                stream = streams[stage.chunk_id % @num_stages]
+                
+                # Sync to ensure transfer complete
+                sync_stream(stream)
+                
+                # Start compute
+                reduce_fn.call(stage.offset, stage.size, stream)
+                stage.state = STATE_COMPUTING
+              end
+            end
+
+            # Check for completed compute
+            active_stages.reject! do |stage|
+              if stage.state == STATE_COMPUTING
+                stream = streams[stage.chunk_id % @num_stages]
+                sync_stream(stream)
+                stage.state = STATE_COMPLETE
+                true
+              else
+                false
+              end
+            end
+          end
+        end
+
+        # Execute a pipelined AllReduce with overlap
+        #
+        # @param ring [Ring] Ring algorithm instance
+        # @param buffers [Array<FFI::Pointer>] Device buffers
+        # @param total_size [Integer] Total buffer size
+        # @param dtype [Symbol] Data type
+        # @param op [Symbol] Reduction operation
+        # @param gpu_streams [Array<Array<CUDA::Stream>>] Streams per GPU per stage
+        # @return [void]
+        def execute_pipelined_all_reduce(ring:, buffers:, total_size:, dtype:, op:, gpu_streams: nil)
+          n_gpus = ring.n_gpus
+          n_chunks = (total_size + @chunk_size - 1) / @chunk_size
+
+          # Create streams: 2 per GPU for double-buffering
+          gpu_streams ||= n_gpus.times.map do |rank|
+            CUDA::RuntimeAPI.cudaSetDevice(ring.ring_order[rank])
+            create_streams_for_device(@num_stages)
+          end
+
+          # Process in pipeline fashion
+          n_chunks.times do |chunk_id|
+            chunk_offset = chunk_id * @chunk_size
+            chunk_size = [total_size - chunk_offset, @chunk_size].min
+
+            stage_idx = chunk_id % @num_stages
+
+            # Get streams for this stage
+            stage_streams = gpu_streams.map { |streams| streams[stage_idx] }
+
+            # Create offset buffers
+            chunk_buffers = buffers.map { |buf| ptr_offset(buf, chunk_offset) }
+
+            # Execute ring all-reduce on this chunk
+            ring.all_reduce(
+              buffers: chunk_buffers,
+              sizes: [chunk_size] * n_gpus,
+              dtype: dtype,
+              op: op,
+              streams: stage_streams
+            )
+          end
+
+          # Final sync
+          gpu_streams.flatten.each { |s| sync_stream(s) }
+        end
+
+        private
+
+        def create_streams(count, buffer)
+          count.times.map do
+            stream_ptr = FFI::MemoryPointer.new(:pointer)
+            CUDA::RuntimeAPI.cudaStreamCreate(stream_ptr)
+            stream_ptr.read_pointer
+          end
+        end
+
+        def create_streams_for_device(count)
+          count.times.map do
+            stream_ptr = FFI::MemoryPointer.new(:pointer)
+            CUDA::RuntimeAPI.cudaStreamCreate(stream_ptr)
+            stream_ptr.read_pointer
+          end
+        end
+
+        def sync_stream(stream)
+          stream_ptr = case stream
+                       when FFI::Pointer then stream
+                       when CUDA::Stream then stream.ptr
+                       else FFI::Pointer::NULL
+                       end
+
+          if stream_ptr.null?
+            CUDA::RuntimeAPI.cudaDeviceSynchronize
+          else
+            CUDA::RuntimeAPI.cudaStreamSynchronize(stream_ptr)
+          end
+        end
+
+        def ptr_offset(ptr, offset)
+          FFI::Pointer.new(:uint8, ptr.address + offset)
+        end
+      end
+    end
+  end
+end