ignis-collective 0.0.1

data/lib/nvruby/collective/communicator.rb

@@ -0,0 +1,633 @@
+# frozen_string_literal: true
+
+require_relative "device_manager"
+require_relative "transport_selector"
+require_relative "topology"
+require_relative "algorithms/reduction_ops"
+require_relative "algorithms/ring"
+require_relative "algorithms/tree"
+
+module Ignis
+  module Collective
+    # Primary user-facing abstraction for collective operations
+    # Provides AllReduce, Broadcast, Reduce, and other collective primitives
+    class Communicator
+      # Reduction operations
+      REDUCTION_OPS = [:sum, :prod, :min, :max, :avg].freeze
+
+      # @return [Array<Integer>] GPU device IDs in this communicator
+      attr_reader :gpu_ids
+
+      # @return [DeviceManager] Device manager
+      attr_reader :device_manager
+
+      # @return [TransportSelector] Transport selector
+      attr_reader :transport_selector
+
+      # @return [Integer] Rank of this communicator (for multi-process)
+      attr_reader :rank
+
+      # @return [Integer] Total number of ranks
+      attr_reader :world_size
+
+      # Create a new communicator for the specified GPUs
+      # @param gpu_ids [Array<Integer>] GPU device IDs to include
+      # @param rank [Integer] Rank of this process (default 0 for single-process)
+      # @param world_size [Integer] Total ranks (default 1 for single-process)
+      def initialize(gpu_ids:, rank: 0, world_size: 1)
+        @gpu_ids = gpu_ids.dup.freeze
+        @rank = rank
+        @world_size = world_size
+
+        validate_gpu_ids!
+
+        @device_manager = DeviceManager.new(device_ids: @gpu_ids)
+        @transport_selector = TransportSelector.new(@gpu_ids)
+        @ring_order = nil
+        @initialized = false
+      end
+
+      # Initialize the communicator (detect topology, enable P2P, etc.)
+      # @return [self]
+      def initialize!
+        return self if @initialized
+
+        @device_manager.initialize!
+        @device_manager.enable_all_p2p_access!
+        @transport_selector.initialize!
+        @ring_order = @transport_selector.optimal_ring_order
+
+        @initialized = true
+        self
+      end
+
+      # Perform AllReduce operation - reduce and distribute result to all GPUs
+      # @param tensors [Array<NvArray>] One tensor per GPU
+      # @param op [Symbol] Reduction operation (:sum, :prod, :min, :max)
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [Array<NvArray>] Reduced tensors (same references as input)
+      def all_reduce(tensors, op: :sum, stream: nil)
+        validate_operation!(op)
+        validate_tensors!(tensors)
+        ensure_initialized!
+
+        # Single GPU case - no-op
+        return tensors if @gpu_ids.size == 1
+
+        # Use Ring AllReduce for multi-GPU
+        ring_all_reduce(tensors, op, stream)
+      end
+
+      # Async AllReduce - requires explicit synchronization
+      # @param tensors [Array<NvArray>] One tensor per GPU
+      # @param op [Symbol] Reduction operation
+      # @param stream [CUDA::Stream] CUDA stream for async execution
+      # @return [Array<NvArray>] Tensors (result available after sync)
+      def all_reduce_async(tensors, op: :sum, stream:)
+        raise ArgumentError, "Stream required for async operation" unless stream
+
+        all_reduce(tensors, op: op, stream: stream)
+      end
+
+      # Broadcast tensor from root GPU to all GPUs
+      # @param tensor [NvArray] Source tensor on root GPU
+      # @param root [Integer] Root GPU index (default 0)
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [Array<NvArray>] Tensors on all GPUs with broadcasted data
+      def broadcast(tensor, root: 0, stream: nil)
+        ensure_initialized!
+        validate_gpu_index!(root)
+
+        return [tensor] if @gpu_ids.size == 1
+
+        # TODO: Implement tree broadcast algorithm
+        # For now, use simple fan-out from root
+        simple_broadcast(tensor, root, stream)
+      end
+
+      # Reduce tensors to root GPU
+      # @param tensors [Array<NvArray>] One tensor per GPU
+      # @param root [Integer] Root GPU index
+      # @param op [Symbol] Reduction operation
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [NvArray] Reduced tensor on root GPU
+      def reduce(tensors, root: 0, op: :sum, stream: nil)
+        validate_operation!(op)
+        validate_tensors!(tensors)
+        ensure_initialized!
+        validate_gpu_index!(root)
+
+        return tensors[0] if @gpu_ids.size == 1
+
+        # TODO: Implement tree reduce algorithm
+        simple_reduce(tensors, root, op, stream)
+      end
+
+      # AllGather - gather tensors from all GPUs to all GPUs
+      # @param tensors [Array<NvArray>] One tensor per GPU (each may be different size)
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [Array<Array<NvArray>>] Gathered tensors on each GPU
+      def all_gather(tensors, stream: nil)
+        validate_tensors!(tensors)
+        ensure_initialized!
+
+        return [tensors] if @gpu_ids.size == 1
+
+        # TODO: Implement ring all-gather
+        simple_all_gather(tensors, stream)
+      end
+
+      # ReduceScatter - reduce and scatter result
+      # @param tensors [Array<NvArray>] One tensor per GPU
+      # @param op [Symbol] Reduction operation
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [Array<FFI::Pointer>] Scattered reduced chunks (chunk size = total_size / N)
+      def reduce_scatter(tensors, op: :sum, stream: nil)
+        validate_operation!(op)
+        validate_tensors!(tensors)
+        ensure_initialized!
+
+        return tensors if @gpu_ids.size == 1
+
+        ring = Algorithms::Ring.new(
+          ring_order: @ring_order,
+          transport_selector: @transport_selector
+        )
+
+        buffers = tensors.map { |t| device_buffer(t) }
+        sizes = tensors.map { |t| byte_size_of(t) }
+
+        dtype = if tensors[0].respond_to?(:dtype)
+                  tensors[0].dtype
+                else
+                  :float32
+                end
+
+        # Calculate chunk size
+        total_size = sizes[0]
+        chunk_size = ring.calculate_chunk_size(total_size)
+
+        # Allocate result buffers
+        result_buffers = @gpu_ids.map do |gpu_id|
+          allocate_buffer_on_device(gpu_id, chunk_size)
+        end
+
+        streams = stream ? [stream] * @gpu_ids.size : create_null_streams(@gpu_ids.size)
+
+        ring.reduce_scatter(
+          buffers: buffers,
+          result_buffers: result_buffers,
+          sizes: sizes,
+          dtype: dtype,
+          op: op,
+          streams: streams
+        )
+
+        result_buffers
+      end
+
+      # AllToAll - full exchange between all GPUs
+      # Each GPU sends N chunks (one to each GPU) and receives N chunks (one from each GPU)
+      # @param send_buffers [Array<Array<FFI::Pointer>>] N×N array: send_buffers[src][dst]
+      # @param recv_buffers [Array<Array<FFI::Pointer>>] N×N array: recv_buffers[dst][src]
+      # @param chunk_size [Integer] Size of each chunk in bytes
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [void]
+      def all_to_all(send_buffers, recv_buffers, chunk_size:, stream: nil)
+        ensure_initialized!
+
+        n = @gpu_ids.size
+        return if n == 1
+
+        streams = stream ? [stream] * n : create_null_streams(n)
+
+        # Phase 1: Copy local data (GPU[i] → GPU[i])
+        n.times do |rank|
+          gpu_id = @gpu_ids[rank]
+          CUDA::RuntimeAPI.cudaSetDevice(gpu_id)
+          stream_ptr = get_stream_ptr(streams[rank])
+
+          CUDA::RuntimeAPI.cudaMemcpyAsync(
+            recv_buffers[rank][rank],
+            send_buffers[rank][rank],
+            chunk_size,
+            CUDA::RuntimeAPI::MEMCPY_DEVICE_TO_DEVICE,
+            stream_ptr
+          )
+        end
+
+        # Phase 2: N-1 rounds of pairwise exchange
+        (n - 1).times do |round|
+          n.times do |rank|
+            gpu_id = @gpu_ids[rank]
+            
+            # Calculate partner for this round (rotation pattern)
+            partner = (rank + round + 1) % n
+            partner_gpu = @gpu_ids[partner]
+
+            stream_ptr = get_stream_ptr(streams[rank])
+
+            # Send to partner
+            transport = @transport_selector.select_transport(gpu_id, partner_gpu)
+            
+            if transport.is_a?(Transport::P2PTransport)
+              transport.copy_async(
+                recv_buffers[partner][rank],  # Partner receives from me
+                send_buffers[rank][partner],  # I send to partner
+                chunk_size,
+                stream_ptr
+              )
+            end
+          end
+
+          # Synchronize after each round
+          synchronize_all_streams!(streams)
+        end
+      end
+
+      # Point-to-point send from current rank to destination
+      # @param tensor [NvArray, FFI::Pointer] Data to send
+      # @param dest_rank [Integer] Destination rank (index in gpu_ids)
+      # @param size [Integer, nil] Size in bytes (inferred from tensor if nil)
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [void]
+      def send(tensor, dest_rank:, size: nil, stream: nil)
+        ensure_initialized!
+        validate_gpu_index!(dest_rank)
+
+        src_rank = 0  # Default sender is rank 0
+        src_gpu = @gpu_ids[src_rank]
+        dst_gpu = @gpu_ids[dest_rank]
+
+        return if src_rank == dest_rank
+
+        buffer = device_buffer(tensor)
+        byte_size = size || byte_size_of(tensor)
+
+        transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+        stream_ptr = stream ? get_stream_ptr(stream) : FFI::Pointer::NULL
+
+        if transport.is_a?(Transport::P2PTransport)
+          # P2P copy requires destination buffer
+          # Assumes tensor has been pre-allocated on dest
+          raise ArgumentError, "P2P send requires pre-allocated recv buffer on dest"
+        end
+      end
+
+      # Point-to-point send from specific source rank
+      # @param buffer [FFI::Pointer] Source buffer on src_rank GPU
+      # @param src_rank [Integer] Source rank
+      # @param dst_buffer [FFI::Pointer] Destination buffer on dst_rank GPU
+      # @param dst_rank [Integer] Destination rank
+      # @param size [Integer] Size in bytes
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [void]
+      def send_recv(buffer, src_rank:, dst_buffer:, dst_rank:, size:, stream: nil)
+        ensure_initialized!
+        validate_gpu_index!(src_rank)
+        validate_gpu_index!(dst_rank)
+
+        return if src_rank == dst_rank
+
+        src_gpu = @gpu_ids[src_rank]
+        dst_gpu = @gpu_ids[dst_rank]
+
+        transport = @transport_selector.select_transport(src_gpu, dst_gpu)
+        stream_ptr = stream ? get_stream_ptr(stream) : FFI::Pointer::NULL
+
+        if transport.is_a?(Transport::P2PTransport)
+          # Set source device context
+          CUDA::RuntimeAPI.cudaSetDevice(src_gpu)
+          transport.copy_async(dst_buffer, buffer, size, stream_ptr)
+        elsif transport.is_a?(Transport::IPCTransport)
+          # For IPC, export/import handles
+          handle = transport.export_handle(buffer)
+          CUDA::RuntimeAPI.cudaSetDevice(dst_gpu)
+          mapped = transport.import_handle(handle)
+          
+          # Copy from mapped to destination
+          CUDA::RuntimeAPI.cudaMemcpyAsync(
+            dst_buffer,
+            mapped,
+            size,
+            CUDA::RuntimeAPI::MEMCPY_DEVICE_TO_DEVICE,
+            stream_ptr
+          )
+          
+          transport.close_imported_handle(mapped)
+        end
+      end
+
+      # Point-to-point receive (no-op, actual receive happens in send_recv)
+      # @param buffer [FFI::Pointer] Buffer to receive into
+      # @param src_rank [Integer] Source rank
+      # @param size [Integer] Expected size in bytes
+      # @param stream [CUDA::Stream, nil] Optional CUDA stream
+      # @return [void]
+      def recv(buffer, src_rank:, size:, stream: nil)
+        ensure_initialized!
+        validate_gpu_index!(src_rank)
+        # Actual data transfer happens via send_recv from sender side
+        # This just marks the receive buffer as ready
+        barrier
+      end
+
+      # Barrier synchronization across all GPUs
+      # @return [void]
+      def barrier
+        ensure_initialized!
+        @device_manager.synchronize_all!
+      end
+
+      # Check if communicator is ready
+      # @return [Boolean] True if initialized
+      def ready?
+        @initialized &&
+          @device_manager.ready? &&
+          @transport_selector.ready?
+      end
+
+      # Get the topology matrix
+      # @return [Topology::Matrix] Topology information
+      def topology
+        @device_manager.topology&.matrix
+      end
+
+      # Get performance summary
+      # @return [Hash] Performance statistics
+      def performance_summary
+        @transport_selector.performance_summary
+      end
+
+      # Clean up all resources
+      # @return [void]
+      def destroy!
+        @transport_selector.destroy!
+        @device_manager.destroy!
+        @initialized = false
+      end
+
+      # @return [String] Human-readable description
+      def to_s
+        status = @initialized ? "ready" : "uninitialized"
+        "Communicator[#{@gpu_ids.size} GPUs, #{status}]"
+      end
+
+      # @return [String] Detailed inspection
+      def inspect
+        "#<Ignis::Collective::Communicator " \
+          "gpu_ids=#{@gpu_ids} " \
+          "rank=#{@rank}/#{@world_size} " \
+          "initialized=#{@initialized}>"
+      end
+
+      private
+
+      def validate_gpu_ids!
+        raise ArgumentError, "gpu_ids cannot be empty" if @gpu_ids.empty?
+
+        max_id = CUDA::Device.count - 1
+        invalid = @gpu_ids.reject { |id| id.between?(0, max_id) }
+        return if invalid.empty?
+
+        raise ArgumentError, "Invalid GPU IDs: #{invalid}. Valid range: 0-#{max_id}"
+      end
+
+      def validate_operation!(op)
+        return if REDUCTION_OPS.include?(op)
+
+        raise ArgumentError, "Invalid reduction op: #{op}. Valid: #{REDUCTION_OPS}"
+      end
+
+      def validate_tensors!(tensors)
+        if tensors.size != @gpu_ids.size
+          raise ArgumentError,
+                "Expected #{@gpu_ids.size} tensors, got #{tensors.size}"
+        end
+      end
+
+      def validate_gpu_index!(index)
+        return if index >= 0 && index < @gpu_ids.size
+
+        raise ArgumentError, "Invalid GPU index: #{index}. Valid: 0-#{@gpu_ids.size - 1}"
+      end
+
+      def ensure_initialized!
+        return if @initialized
+
+        raise CommunicatorError, "Communicator not initialized. Call initialize! first."
+      end
+
+      # Ring AllReduce using the Ring algorithm
+      # Uses scatter-reduce + allgather pattern
+      def ring_all_reduce(tensors, op, stream)
+        n = @gpu_ids.size
+        return tensors if n == 1
+
+        # Create Ring algorithm instance
+        ring = Algorithms::Ring.new(
+          ring_order: @ring_order,
+          transport_selector: @transport_selector
+        )
+
+        # Extract device pointers and BYTE sizes from tensors.
+        buffers = tensors.map { |t| device_buffer(t) }
+        sizes = tensors.map { |t| byte_size_of(t) }
+
+        # Detect dtype (default to float32)
+        dtype = tensors[0].respond_to?(:dtype) ? tensors[0].dtype : :float32
+
+        # Create streams for each GPU
+        streams = stream ? [stream] * n : create_null_streams(n)
+
+        # Execute Ring AllReduce (scatter-reduce + all-gather). For :avg the ring
+        # accumulates a SUM; we divide by the participant count once at the end.
+        ring.all_reduce(
+          buffers: buffers,
+          sizes: sizes,
+          dtype: dtype,
+          op: op,
+          streams: streams
+        )
+
+        apply_avg!(buffers, sizes, dtype, n) if op == :avg
+
+        tensors
+      end
+
+      # Create null stream pointers for each GPU
+      def create_null_streams(n)
+        Array.new(n) { FFI::Pointer::NULL }
+      end
+
+      # Extract an FFI device pointer from a tensor/NvArray (either class) or
+      # pass through a raw pointer. (Was `t.data_ptr` — which no NvArray defines.)
+      def device_buffer(t)
+        if t.respond_to?(:device_ffi_ptr) then t.device_ffi_ptr   # Ignis/Shared NvArray
+        elsif t.respond_to?(:device_ptr) then t.device_ptr
+        elsif t.respond_to?(:data_ptr) then t.data_ptr
+        else t                                                     # assume FFI::Pointer
+        end
+      end
+
+      # Size of a tensor in BYTES. (Was `t.byte_size || t.size`, and t.size is the
+      # element COUNT — not bytes — for NvArray, so reductions ran on 1/4 the data.)
+      def byte_size_of(t)
+        if t.respond_to?(:nbytes) then t.nbytes          # Ignis::NvArray
+        elsif t.respond_to?(:size_bytes) then t.size_bytes # Ignis::Shared::NvArray
+        elsif t.respond_to?(:byte_size) then t.byte_size
+        else raise ArgumentError, "cannot determine byte size of #{t.class}"
+        end
+      end
+
+      # Element size in bytes for a dtype.
+      def elem_size_of(dtype)
+        case dtype
+        when :float64, :int64 then 8
+        when :float16, :bfloat16 then 2
+        when :int8, :uint8 then 1
+        else 4
+        end
+      end
+
+      # Divide each buffer by the participant count in place (for op: :avg).
+      def apply_avg!(buffers, sizes, dtype, n)
+        scale = 1.0 / n
+        kernel = Ignis::JIT::Kernels::Elementwise.scale_forward
+        es = elem_size_of(dtype)
+        buffers.each_with_index do |buf, i|
+          count = sizes[i] / es
+          CUDA::RuntimeAPI.cudaSetDevice(@gpu_ids[i])
+          kernel.launch(grid: [(count + 255) / 256], block: [256], args: [buf, buf, scale, count])
+        end
+        Ignis.synchronize
+      end
+
+      # Tree broadcast from root
+      def simple_broadcast(tensor, root, stream)
+        # Create buffers array with tensor at root position
+        buffers = @gpu_ids.map.with_index do |gpu_id, i|
+          if i == root
+            device_buffer(tensor)
+          else
+            # Allocate buffer on other GPUs
+            allocate_buffer_on_device(gpu_id, byte_size_of(tensor))
+          end
+        end
+
+        size = byte_size_of(tensor)
+
+        # Create Tree algorithm instance
+        tree = Algorithms::Tree.new(
+          gpu_ids: @gpu_ids,
+          transport_selector: @transport_selector
+        )
+
+        # Create streams
+        streams = stream ? [stream] * @gpu_ids.size : create_null_streams(@gpu_ids.size)
+
+        # Execute broadcast
+        tree.broadcast(
+          buffer: buffers[root],
+          buffers: buffers,
+          size: size,
+          root: root,
+          streams: streams
+        )
+
+        buffers
+      end
+
+      # Tree reduce to root
+      def simple_reduce(tensors, root, op, stream)
+        buffers = tensors.map { |t| device_buffer(t) }
+        sizes = tensors.map { |t| byte_size_of(t) }
+
+        dtype = if tensors[0].respond_to?(:dtype)
+                  tensors[0].dtype
+                else
+                  :float32
+                end
+
+        tree = Algorithms::Tree.new(
+          gpu_ids: @gpu_ids,
+          transport_selector: @transport_selector
+        )
+
+        streams = stream ? [stream] * @gpu_ids.size : create_null_streams(@gpu_ids.size)
+
+        tree.reduce(
+          buffers: buffers,
+          sizes: sizes,
+          dtype: dtype,
+          op: op,
+          root: root,
+          streams: streams
+        )
+
+        tensors[root]
+      end
+
+      # All-gather. In this single-process model every rank's tensor is already
+      # an accessible Ruby/GPU object, so the gathered set on each rank is simply
+      # all input tensors. We return an independent list per rank (so callers may
+      # mutate one rank's view without aliasing others).
+      #
+      # NOTE: true cross-PROCESS / cross-node gather (world_size > 1) requires the
+      # transport layer to physically move device buffers, which is not wired yet;
+      # ring AllGather machinery exists in Algorithms::Ring#all_gather_standalone
+      # for when P2P/host-staged transports are completed.
+      def simple_all_gather(tensors, _stream)
+        barrier
+        Array.new(@gpu_ids.size) { tensors.dup }
+      end
+
+      # Helper: get tensor size
+      def tensor_size(tensor)
+        if tensor.respond_to?(:byte_size)
+          tensor.byte_size
+        elsif tensor.respond_to?(:size)
+          tensor.size
+        else
+          4  # Default to 4 bytes
+        end
+      end
+
+      # Helper: allocate buffer on specific device
+      def allocate_buffer_on_device(gpu_id, size)
+        CUDA::RuntimeAPI.ensure_loaded!
+        CUDA::RuntimeAPI.cudaSetDevice(gpu_id)
+        ptr_ptr = FFI::MemoryPointer.new(:pointer)
+        status = CUDA::RuntimeAPI.cudaMalloc(ptr_ptr, size)
+        CUDA::RuntimeAPI.check_status!(status, "Alloc broadcast buffer")
+        ptr_ptr.read_pointer
+      end
+
+      # Helper: get stream pointer for FFI
+      def get_stream_ptr(stream)
+        case stream
+        when FFI::Pointer
+          stream
+        when CUDA::Stream
+          stream.ptr
+        else
+          FFI::Pointer::NULL
+        end
+      end
+
+      # Helper: synchronize all streams
+      def synchronize_all_streams!(streams)
+        streams.each_with_index do |stream, i|
+          CUDA::RuntimeAPI.cudaSetDevice(@gpu_ids[i])
+          
+          stream_ptr = get_stream_ptr(stream)
+          if stream_ptr.null?
+            CUDA::RuntimeAPI.cudaDeviceSynchronize
+          else
+            CUDA::RuntimeAPI.cudaStreamSynchronize(stream_ptr)
+          end
+        end
+      end
+    end
+  end
+end