ignis-collective 0.0.1

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

@@ -0,0 +1,284 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    module Algorithms
+      # Topology-Aware Routing Optimizer
+      #
+      # Optimizes collective operation routing based on PCIe/NVLink topology.
+      # Groups GPUs by PCIe switch to minimize cross-switch traffic.
+      #
+      # Key optimizations:
+      # 1. Intra-switch communication first (lower latency)
+      # 2. Ring order follows physical topology
+      # 3. Tree structure respects switch boundaries
+      class TopologyRouter
+        # Switch group for routing
+        SwitchGroup = Struct.new(:switch_id, :gpu_ids, :bandwidth_gbps, keyword_init: true)
+
+        # @return [Topology::Detector] Topology detector
+        attr_reader :topology
+
+        # @return [Array<Integer>] GPU IDs
+        attr_reader :gpu_ids
+
+        # @return [Array<SwitchGroup>] GPU groups by switch
+        attr_reader :switch_groups
+
+        # @param gpu_ids [Array<Integer>] GPU device IDs
+        # @param topology [Topology::Detector, nil] Optional topology detector
+        def initialize(gpu_ids:, topology: nil)
+          @gpu_ids = gpu_ids.dup.freeze
+          @topology = topology || Topology::Detector.new
+          @switch_groups = []
+          @routing_table = {}
+          detect_switch_groups!
+        end
+
+        # Get optimized ring order based on topology
+        # Groups GPUs by PCIe switch, then chains switches
+        #
+        # @return [Array<Integer>] Optimal ring order
+        def optimal_ring_order
+          return @gpu_ids if @switch_groups.size <= 1
+
+          # Build ring by chaining switch groups
+          # Within each group, order by P2P performance
+          ring = []
+
+          # Sort groups by average bandwidth (best first)
+          sorted_groups = @switch_groups.sort_by { |g| -g.bandwidth_gbps }
+
+          sorted_groups.each do |group|
+            # Order GPUs within group by P2P performance
+            ordered = order_within_group(group.gpu_ids)
+            ring.concat(ordered)
+          end
+
+          ring
+        end
+
+        # Get optimized tree structure that respects switch boundaries
+        # 
+        # @param root [Integer] Root rank
+        # @return [Hash] Tree structure with :parent, :children for each rank
+        def optimal_tree_structure(root: 0)
+          tree = {}
+          n = @gpu_ids.size
+
+          return { 0 => { parent: nil, children: [] } } if n == 1
+
+          # Build hierarchical tree:
+          # - Level 0: Root
+          # - Level 1: One representative from each switch group
+          # - Level 2+: Other GPUs within each switch group
+
+          # Find which group the root belongs to
+          root_group_idx = @switch_groups.find_index { |g| g.gpu_ids.include?(@gpu_ids[root]) }
+          root_group = @switch_groups[root_group_idx]
+
+          # Root has no parent
+          tree[root] = { parent: nil, children: [] }
+
+          # Connect switch group representatives to root
+          group_reps = []
+          @switch_groups.each_with_index do |group, idx|
+            next if idx == root_group_idx
+
+            # Pick highest-bandwidth GPU in group as representative
+            rep_gpu = pick_best_connector(root_group, group)
+            rep_rank = @gpu_ids.index(rep_gpu)
+            
+            tree[root][:children] << rep_rank
+            tree[rep_rank] = { parent: root, children: [] }
+            group_reps << { rank: rep_rank, group: group }
+          end
+
+          # Connect other GPUs in root's group directly to root
+          root_group.gpu_ids.each do |gpu|
+            rank = @gpu_ids.index(gpu)
+            next if rank == root
+            
+            tree[root][:children] << rank
+            tree[rank] = { parent: root, children: [] }
+          end
+
+          # Connect remaining GPUs in each group to their representative
+          group_reps.each do |rep_info|
+            rep_rank = rep_info[:rank]
+            rep_gpu = @gpu_ids[rep_rank]
+            
+            rep_info[:group].gpu_ids.each do |gpu|
+              next if gpu == rep_gpu
+              
+              rank = @gpu_ids.index(gpu)
+              tree[rep_rank][:children] << rank
+              tree[rank] = { parent: rep_rank, children: [] }
+            end
+          end
+
+          tree
+        end
+
+        # Suggest optimal algorithm based on message size and topology
+        #
+        # @param message_size [Integer] Total message size in bytes
+        # @return [Symbol] :ring, :tree, or :double_tree
+        def suggest_algorithm(message_size)
+          # Ring is bandwidth-optimal for large messages
+          # Tree is latency-optimal for small messages
+          # Double tree for non-power-of-2 with medium messages
+
+          n = @gpu_ids.size
+          is_power_of_2 = n > 0 && (n & (n - 1)).zero?
+
+          if message_size < 1024  # < 1KB
+            :tree
+          elsif message_size > 1_048_576  # > 1MB
+            :ring
+          elsif !is_power_of_2
+            :double_tree
+          else
+            # Medium messages: prefer ring if good interconnect
+            avg_bandwidth = @switch_groups.sum { |g| g.bandwidth_gbps } / @switch_groups.size.to_f
+            avg_bandwidth > 50 ? :ring : :tree
+          end
+        end
+
+        # Get routing path between two GPUs
+        #
+        # @param src_rank [Integer] Source rank
+        # @param dst_rank [Integer] Destination rank
+        # @return [Array<Integer>] Path as list of rank hops
+        def get_path(src_rank, dst_rank)
+          return [src_rank] if src_rank == dst_rank
+
+          src_gpu = @gpu_ids[src_rank]
+          dst_gpu = @gpu_ids[dst_rank]
+
+          src_group = find_group(src_gpu)
+          dst_group = find_group(dst_gpu)
+
+          if src_group == dst_group
+            # Direct path within same switch
+            [src_rank, dst_rank]
+          else
+            # Need to cross switch boundary - use group representative
+            # Find best intermediate hop
+            intermediate = pick_best_connector_rank(src_group, dst_group)
+            
+            if intermediate == src_rank || intermediate == dst_rank
+              [src_rank, dst_rank]
+            else
+              [src_rank, intermediate, dst_rank]
+            end
+          end
+        end
+
+        private
+
+        def detect_switch_groups!
+          # Use topology to group GPUs by PCIe switch
+          # GPUs with NVLink form one group, others grouped by P2P connectivity
+
+          matrix = @topology.matrix
+          visited = Set.new
+          group_id = 0
+
+          @gpu_ids.each do |gpu|
+            next if visited.include?(gpu)
+
+            group_gpus = [gpu]
+            visited.add(gpu)
+
+            # Find all GPUs that have NVLink or high-speed P2P to this one
+            @gpu_ids.each do |other|
+              next if visited.include?(other) || other == gpu
+
+              path = matrix.path(gpu, other)
+              if path && (path.nvlink? || path.pcie_p2p?)
+                # Check if they're likely on same switch (both have good P2P to each other)
+                reverse_path = matrix.path(other, gpu)
+                if reverse_path && (reverse_path.nvlink? || reverse_path.pcie_p2p?)
+                  group_gpus << other
+                  visited.add(other)
+                end
+              end
+            end
+
+            # Calculate average bandwidth for this group
+            total_bw = 0.0
+            count = 0
+            group_gpus.combination(2).each do |a, b|
+              path = matrix.path(a, b)
+              if path
+                total_bw += path.bandwidth_gbps
+                count += 1
+              end
+            end
+            avg_bw = count > 0 ? total_bw / count : 16.0  # Default PCIe
+
+            @switch_groups << SwitchGroup.new(
+              switch_id: group_id,
+              gpu_ids: group_gpus,
+              bandwidth_gbps: avg_bw
+            )
+            group_id += 1
+          end
+        end
+
+        def order_within_group(gpus)
+          return gpus if gpus.size <= 2
+
+          # Order by greedy nearest-neighbor based on P2P performance
+          matrix = @topology.matrix
+          remaining = gpus.dup
+          ordered = [remaining.shift]
+
+          while remaining.any?
+            current = ordered.last
+            best = remaining.max_by do |gpu|
+              path = matrix.path(current, gpu)
+              path ? path.performance_rank : 0
+            end
+            ordered << best
+            remaining.delete(best)
+          end
+
+          ordered
+        end
+
+        def pick_best_connector(group_a, group_b)
+          # Find GPU in group_b that has best connection to group_a
+          matrix = @topology.matrix
+          
+          best_gpu = nil
+          best_score = -1
+
+          group_b.gpu_ids.each do |gpu_b|
+            score = group_a.gpu_ids.sum do |gpu_a|
+              path = matrix.path(gpu_a, gpu_b)
+              path ? path.performance_rank : 0
+            end
+            
+            if score > best_score
+              best_score = score
+              best_gpu = gpu_b
+            end
+          end
+
+          best_gpu || group_b.gpu_ids.first
+        end
+
+        def pick_best_connector_rank(group_a, group_b)
+          gpu = pick_best_connector(group_a, group_b)
+          @gpu_ids.index(gpu)
+        end
+
+        def find_group(gpu)
+          @switch_groups.find { |g| g.gpu_ids.include?(gpu) }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,291 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    module Algorithms
+      # Tree-based algorithms for broadcast and reduce operations
+      #
+      # Binary tree algorithms are optimal for small messages where latency dominates:
+      # - Latency complexity: O(log N) steps
+      # - Bandwidth complexity: O(data_size) per step (not optimal for large messages)
+      #
+      # Best for: Small messages (<1KB) where latency is critical
+      class Tree
+        # Tree node representing a GPU in the tree structure
+        TreeNode = Struct.new(:gpu_id, :rank, :parent, :children, :depth, 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
+
+        # @return [Array<TreeNode>] Tree structure
+        attr_reader :tree
+
+        # @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
+          @tree = build_tree
+        end
+
+        # Broadcast data from root to all GPUs
+        #
+        # @param buffer [FFI::Pointer] Source buffer on root GPU
+        # @param buffers [Array<FFI::Pointer>] Destination buffers on all GPUs
+        # @param size [Integer] Buffer size in bytes
+        # @param root [Integer] Root rank (index in gpu_ids)
+        # @param streams [Array<CUDA::Stream, FFI::Pointer>] CUDA streams
+        # @return [void]
+        def broadcast(buffer:, buffers:, size:, root:, streams:)
+          return if @n_gpus == 1
+
+          validate_root!(root)
+
+          # Build tree rooted at specified root
+          tree = build_tree(root: root)
+
+          # Broadcast down the tree: log2(N) steps
+          depth = tree_depth(tree)
+
+          depth.times do |d|
+            # At each depth level, nodes at depth d send to their children
+            tree.each do |node|
+              next unless node.depth == d
+              next if node.children.empty?
+
+              src_rank = node.rank
+              src_gpu = node.gpu_id
+              src_buffer = buffers[src_rank]
+              stream_ptr = get_stream_ptr(streams[src_rank])
+
+              node.children.each do |child_rank|
+                dst_gpu = @gpu_ids[child_rank]
+                dst_buffer = buffers[child_rank]
+
+                transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+                move!(transport, dst_buffer, src_buffer, size, stream_ptr)
+              end
+            end
+
+            # Synchronize after each level
+            synchronize_all_streams!(streams)
+          end
+        end
+
+        # Reduce data from all GPUs to root
+        #
+        # @param buffers [Array<FFI::Pointer>] Source buffers on all GPUs
+        # @param sizes [Array<Integer>] Buffer sizes
+        # @param dtype [Symbol] Data type
+        # @param op [Symbol] Reduction operation
+        # @param root [Integer] Root rank
+        # @param streams [Array<CUDA::Stream, FFI::Pointer>] CUDA streams
+        # @return [void]
+        def reduce(buffers:, sizes:, dtype:, op:, root:, streams:)
+          return if @n_gpus == 1
+
+          validate_root!(root)
+
+          size = sizes[0]
+          elem_size = dtype_elem_size(dtype)
+          elem_count = size / elem_size
+
+          # Build tree rooted at root
+          tree = build_tree(root: root)
+
+          # Allocate temp buffers
+          recv_buffers = allocate_recv_buffers(size)
+
+          begin
+            # Reduce up the tree: log2(N) steps
+            depth = tree_depth(tree)
+
+            # Start from leaves, work up to root
+            (depth - 1).downto(0) do |d|
+              tree.each do |node|
+                next unless node.depth == d
+                next if node.children.empty?
+
+                dst_rank = node.rank
+                dst_gpu = node.gpu_id
+                dst_buffer = buffers[dst_rank]
+                stream_ptr = get_stream_ptr(streams[dst_rank])
+
+                # Set device for reduction
+                CUDA::RuntimeAPI.cudaSetDevice(dst_gpu)
+
+                # For each child: copy its buffer into recv, then reduce it into
+                # the parent's local buffer BEFORE the next child overwrites recv.
+                # (Previously every child was copied into the SAME recv buffer and
+                # then reduced once per child, which double-counted the last child
+                # and dropped the rest — e.g. parent += child2 + child2.)
+                recv_buffer = recv_buffers[dst_rank]
+                node.children.each do |child_rank|
+                  src_gpu = @gpu_ids[child_rank]
+                  src_buffer = buffers[child_rank]
+
+                  transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+                  move!(transport, recv_buffer, src_buffer, size, stream_ptr)
+                  synchronize_stream!(streams[dst_rank])
+
+                  ReductionOps.execute(op, dst_buffer, recv_buffer, dst_buffer,
+                                       elem_count, dtype, stream_ptr)
+                  synchronize_stream!(streams[dst_rank])
+                end
+              end
+
+              synchronize_all_streams!(streams)
+            end
+          ensure
+            free_recv_buffers(recv_buffers)
+          end
+        end
+
+        private
+
+        # Move bytes via the selected transport, failing LOUDLY rather than
+        # silently skipping. A non-P2P transport without copy_async would
+        # otherwise drop data and corrupt the reduction with no error.
+        def move!(transport, dst, src, size, stream_ptr)
+          if transport.respond_to?(:copy_async)
+            transport.copy_async(dst, src, size, stream_ptr)
+          else
+            raise NotImplementedError,
+                  "Transport #{transport.class} has no copy_async; non-P2P tree " \
+                  "movement is not wired yet (refusing to silently drop data)"
+          end
+        end
+
+        # Build binary tree structure
+        # @param root [Integer] Root rank (default 0)
+        # @return [Array<TreeNode>] Tree nodes
+        def build_tree(root: 0)
+          nodes = @n_gpus.times.map do |rank|
+            TreeNode.new(
+              gpu_id: @gpu_ids[rank],
+              rank: rank,
+              parent: nil,
+              children: [],
+              depth: 0
+            )
+          end
+
+          # Build binary tree with specified root
+          # Use BFS to assign depths and parent/child relationships
+          if root != 0
+            # Reorder so root is first in tree building
+            reordered = [root] + (0...@n_gpus).reject { |i| i == root }
+          else
+            reordered = (0...@n_gpus).to_a
+          end
+
+          # Build tree using heap-like structure
+          reordered.each_with_index do |rank, idx|
+            next if idx.zero?  # Root has no parent
+
+            parent_idx = (idx - 1) / 2
+            parent_rank = reordered[parent_idx]
+
+            nodes[rank].parent = parent_rank
+            nodes[parent_rank].children << rank
+          end
+
+          # Calculate depths via BFS
+          queue = [root]
+          nodes[root].depth = 0
+
+          until queue.empty?
+            current_rank = queue.shift
+            current_depth = nodes[current_rank].depth
+
+            nodes[current_rank].children.each do |child_rank|
+              nodes[child_rank].depth = current_depth + 1
+              queue << child_rank
+            end
+          end
+
+          nodes
+        end
+
+        # Get maximum tree depth
+        def tree_depth(tree)
+          tree.map(&:depth).max + 1
+        end
+
+        def validate_root!(root)
+          unless root >= 0 && root < @n_gpus
+            raise ArgumentError, "Invalid root #{root}. Valid: 0-#{@n_gpus - 1}"
+          end
+        end
+
+        # Allocate temp buffers
+        def allocate_recv_buffers(size)
+          CUDA::RuntimeAPI.ensure_loaded!
+
+          @gpu_ids.map do |gpu_id|
+            CUDA::RuntimeAPI.cudaSetDevice(gpu_id)
+            ptr_ptr = FFI::MemoryPointer.new(:pointer)
+            status = CUDA::RuntimeAPI.cudaMalloc(ptr_ptr, size)
+            CUDA::RuntimeAPI.check_status!(status, "Alloc tree recv buffer")
+            ptr_ptr.read_pointer
+          end
+        end
+
+        # Free temp buffers
+        def free_recv_buffers(recv_buffers)
+          recv_buffers.each_with_index do |buf, i|
+            next unless buf && !buf.null?
+
+            CUDA::RuntimeAPI.cudaSetDevice(@gpu_ids[i])
+            CUDA::RuntimeAPI.cudaFree(buf)
+          rescue StandardError
+            # Ignore cleanup errors
+          end
+        end
+
+        def get_stream_ptr(stream)
+          case stream
+          when FFI::Pointer
+            stream
+          when CUDA::Stream
+            stream.ptr
+          else
+            FFI::Pointer::NULL
+          end
+        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 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