ignis-collective 0.0.1

data/lib/nvruby/collective/array_ops.rb

@@ -0,0 +1,240 @@
+# frozen_string_literal: true
+
+module Ignis
+  module Collective
+    # NvArray Integration for Collective Operations
+    #
+    # Provides tensor-aware wrappers around collective primitives
+    # that work directly with NvArray tensors.
+    module ArrayOps
+      class << self
+        # AllReduce on NvArray tensors
+        #
+        # @param tensors [Array<NvArray>] One NvArray per GPU
+        # @param op [Symbol] Reduction operation (:sum, :prod, :min, :max)
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [Array<NvArray>] Modified tensors with reduced values
+        def all_reduce(tensors, op: :sum, comm:, stream: nil)
+          validate_tensors!(tensors)
+          validate_devices!(tensors, comm)
+
+          comm.all_reduce(tensors, op: op, stream: stream)
+        end
+
+        # Broadcast NvArray from root to all GPUs
+        #
+        # @param tensor [NvArray] Tensor on root GPU
+        # @param root [Integer] Root rank
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [Array<NvArray>] NvArrays on all GPUs with same data
+        def broadcast(tensor, root: 0, comm:, stream: nil)
+          validate_single_tensor!(tensor)
+
+          # Comm.broadcast returns FFI::Pointers, wrap in NvArrays
+          pointers = comm.broadcast(tensor, root: root, stream: stream)
+
+          pointers.map.with_index do |ptr, i|
+            if i == root
+              tensor
+            else
+              create_nvarray_from_ptr(
+                ptr: ptr,
+                shape: tensor.shape,
+                dtype: tensor.dtype,
+                device: comm.gpu_ids[i]
+              )
+            end
+          end
+        end
+
+        # Reduce to root GPU
+        #
+        # @param tensors [Array<NvArray>] One NvArray per GPU
+        # @param root [Integer] Root rank
+        # @param op [Symbol] Reduction operation
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [NvArray] Reduced tensor on root GPU
+        def reduce(tensors, root: 0, op: :sum, comm:, stream: nil)
+          validate_tensors!(tensors)
+
+          result = comm.reduce(tensors, root: root, op: op, stream: stream)
+
+          # result is already the root tensor or pointer
+          result.is_a?(NvArray) ? result : tensors[root]
+        end
+
+        # AllGather - each GPU contributes, all receive all
+        #
+        # @param tensors [Array<NvArray>] One NvArray per GPU (local chunks)
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [Array<NvArray>] NvArrays on all GPUs with gathered data
+        def all_gather(tensors, comm:, stream: nil)
+          validate_tensors!(tensors)
+
+          # Calculate total size after gathering
+          n_gpus = tensors.size
+          chunk_shape = tensors[0].shape.dup
+          gathered_shape = chunk_shape.dup
+          gathered_shape[0] = chunk_shape[0] * n_gpus  # Concatenate along first axis
+
+          gathered_pointers = comm.all_gather(tensors, stream: stream)
+
+          # Wrap results in NvArrays
+          gathered_pointers.map.with_index do |ptr_or_arr, i|
+            if ptr_or_arr.is_a?(Array)
+              # Already NvArrays
+              ptr_or_arr[i]
+            else
+              create_nvarray_from_ptr(
+                ptr: ptr_or_arr,
+                shape: gathered_shape,
+                dtype: tensors[0].dtype,
+                device: comm.gpu_ids[i]
+              )
+            end
+          end
+        end
+
+        # ReduceScatter - reduce then scatter
+        #
+        # @param tensors [Array<NvArray>] One NvArray per GPU (full size)
+        # @param op [Symbol] Reduction operation
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [Array<NvArray>] Chunk NvArrays (1/N size each)
+        def reduce_scatter(tensors, op: :sum, comm:, stream: nil)
+          validate_tensors!(tensors)
+
+          n_gpus = tensors.size
+          full_shape = tensors[0].shape.dup
+          chunk_shape = full_shape.dup
+          chunk_shape[0] = full_shape[0] / n_gpus
+
+          chunk_pointers = comm.reduce_scatter(tensors, op: op, stream: stream)
+
+          chunk_pointers.map.with_index do |ptr, i|
+            create_nvarray_from_ptr(
+              ptr: ptr,
+              shape: chunk_shape,
+              dtype: tensors[0].dtype,
+              device: comm.gpu_ids[i]
+            )
+          end
+        end
+
+        # Distribute tensor across GPUs (scatter)
+        #
+        # @param tensor [NvArray] Source tensor on root GPU
+        # @param root [Integer] Root rank
+        # @param comm [Communicator] Communicator instance
+        # @param stream [CUDA::Stream, nil] Optional CUDA stream
+        # @return [Array<NvArray>] Chunks on each GPU
+        def scatter(tensor, root: 0, comm:, stream: nil)
+          validate_single_tensor!(tensor)
+
+          n_gpus = comm.gpu_ids.size
+          full_shape = tensor.shape.dup
+          chunk_shape = full_shape.dup
+          chunk_shape[0] = full_shape[0] / n_gpus
+          chunk_size = tensor.nbytes / n_gpus
+
+          chunks = []
+
+          n_gpus.times do |i|
+            if i == root
+              # Root keeps first chunk
+              chunk_ptr = tensor.device_ptr
+              chunks << create_nvarray_from_ptr(
+                ptr: chunk_ptr,
+                shape: chunk_shape,
+                dtype: tensor.dtype,
+                device: comm.gpu_ids[i]
+              )
+            else
+              # Allocate and send chunk
+              chunk_ptr = allocate_on_device(comm.gpu_ids[i], chunk_size)
+              offset = i * chunk_size
+              src_ptr = FFI::Pointer.new(:uint8, tensor.device_ptr.address + offset)
+
+              comm.send_recv(
+                src_ptr,
+                src_rank: root,
+                dst_buffer: chunk_ptr,
+                dst_rank: i,
+                size: chunk_size,
+                stream: stream
+              )
+
+              chunks << create_nvarray_from_ptr(
+                ptr: chunk_ptr,
+                shape: chunk_shape,
+                dtype: tensor.dtype,
+                device: comm.gpu_ids[i]
+              )
+            end
+          end
+
+          chunks
+        end
+
+        private
+
+        def validate_tensors!(tensors)
+          raise ArgumentError, "Expected array of tensors" unless tensors.is_a?(Array)
+          raise ArgumentError, "Empty tensor array" if tensors.empty?
+
+          first = tensors[0]
+          raise ArgumentError, "Tensors must be NvArray" unless first.is_a?(NvArray)
+          raise ArgumentError, "Tensors must be on device" unless first.on_device?
+
+          tensors.each do |t|
+            unless t.shape == first.shape
+              raise ArgumentError, "All tensors must have same shape"
+            end
+            unless t.dtype == first.dtype
+              raise ArgumentError, "All tensors must have same dtype"
+            end
+          end
+        end
+
+        def validate_single_tensor!(tensor)
+          raise ArgumentError, "Expected NvArray" unless tensor.is_a?(NvArray)
+          raise ArgumentError, "Tensor must be on device" unless tensor.on_device?
+        end
+
+        def validate_devices!(tensors, comm)
+          tensors.each_with_index do |t, i|
+            expected_device = comm.gpu_ids[i]
+            actual_device = t.device
+
+            unless actual_device == expected_device
+              raise ArgumentError, "Tensor #{i} on device #{actual_device}, expected #{expected_device}"
+            end
+          end
+        end
+
+        def create_nvarray_from_ptr(ptr:, shape:, dtype:, device:)
+          NvArray.from_device_ptr(
+            ptr,
+            shape: shape,
+            dtype: dtype,
+            take_ownership: true
+          )
+        end
+
+        def allocate_on_device(device_id, size)
+          CUDA::RuntimeAPI.ensure_loaded!
+          CUDA::RuntimeAPI.cudaSetDevice(device_id)
+          ptr_ptr = FFI::MemoryPointer.new(:pointer)
+          status = CUDA::RuntimeAPI.cudaMalloc(ptr_ptr, size)
+          CUDA::RuntimeAPI.check_status!(status, "Alloc for scatter")
+          ptr_ptr.read_pointer
+        end
+      end
+    end
+  end
+end