Adding Read Put Packet operation at Executor

python/mscclpp/language/dag/instruction_dag.py

@@ -6,6 +6,7 @@
 from mscclpp.language.buffer import Buffer
 from mscclpp.language.types import (
     Channel,
+    DataFormat,
     ChannelType,
     ChunkRef,
     Instruction,
@@ -51,7 +52,7 @@ def convert_set_list(self):
         return visited
 
     def complete_channels(self):
-        send_op = [Instruction.put, Instruction.signal, Instruction.put_packet]
+        send_op = [Instruction.put, Instruction.signal, Instruction.put_packet, Instruction.read_put_packet]
         recv_op = [Instruction.wait, Instruction.get, Instruction.read_reduce_copy]
         group_send_op = [Instruction.group_store]
         group_recv_op = [Instruction.group_load_reduce]
@@ -140,20 +141,33 @@ def add_reduce(self, rank, send_ref, recv_ref, tb, use_packet=False):
         return op
 
     # InstructionDAG - adds a put node
-    def add_put(self, rank, send_ref, recv_ref, tb, ch_type, use_packet=False):
+    def add_put(self, rank, send_ref, recv_ref, tb, src_format, ch_type, use_packet=False):
         tb_step = self._get_tb_step(rank, tb)
         if use_packet:
-            op = Op(
-                Instruction.put_packet,
-                rank,
-                send_ref,
-                recv_ref,
-                next=set(),
-                prev=set(),
-                tb=tb,
-                channel_type=ch_type,
-                step=tb_step,
-            )
+            if src_format == DataFormat.raw:
+                op = Op(
+                    Instruction.put_packet,
+                    rank,
+                    send_ref,
+                    recv_ref,
+                    next=set(),
+                    prev=set(),
+                    tb=tb,
+                    channel_type=ch_type,
+                    step=tb_step,
+                )
+            elif src_format == DataFormat.packet:
+                op = Op(
+                    Instruction.read_put_packet,
+                    rank,
+                    send_ref,
+                    recv_ref,
+                    next=set(),
+                    prev=set(),
+                    tb=tb,
+                    channel_type=ch_type,
+                    step=tb_step,
+                )
         else:
             op = Op(
                 Instruction.put,

python/mscclpp/language/ir.py

@@ -12,6 +12,7 @@
 _local_src_insts_mscclpp: set = {
     Instruction.put,
     Instruction.put_packet,
+    Instruction.read_put_packet,
     Instruction.signal,
     Instruction.flush,
     Instruction.put_with_signal,
@@ -428,6 +429,7 @@ def to_json(self, op: Op, tb_channel_dict: dict) -> _JsonInstruction:
     Instruction.barrier: _BarrierConverter(),
     Instruction.put: _PutConverter(),
     Instruction.put_packet: _PutConverter(),
+    Instruction.read_put_packet: _PutConverter(),
     Instruction.put_with_signal: _PutConverter(),
     Instruction.put_with_signal_and_flush: _PutConverter(),
     Instruction.get: _GetConverter(),

python/mscclpp/language/program.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from mscclpp.language.collectives import Collective
 from mscclpp.language.buffer import *
-from mscclpp.language.types import ChannelType, ChunkRef, ReplicationPolicy, Threadblock
+from mscclpp.language.types import DataFormat, ChannelType, ChunkRef, ReplicationPolicy, Threadblock
 from mscclpp.language.ir import *
 from mscclpp.language.dag import DagOptimizer, DagLower, InstructionDAG
 from mscclpp.language.rank import Rank
@@ -223,41 +223,51 @@ def _get_buffer_index(self, remote_rank, buffer, index):
             return buffer, self.prog.buffers[remote_rank][buffer].instance_size()
         return buffer, index
 
-    def _put(self, dst, buffer=None, index=-1, sendtb=-1, chan_type=ChannelType.memory, use_packet=False):
+    def _put(
+        self,
+        dst,
+        buffer=None,
+        index=-1,
+        sendtb=-1,
+        src_format=DataFormat.raw,
+        chan_type=ChannelType.memory,
+        use_packet=False,
+    ):
         self.prog.check_buffer_exists(dst, buffer)
         assert self.rank != dst, "Cannot put to the same rank"
         buffer, index = self._get_buffer_index(dst, buffer, index)
 
         dst_chunkref = self.prog.get_ref(dst, buffer, index, self.size)
         self.prog.apply_send(self.rank, self.buffer, self.index, dst, buffer, index, self.size)
         if use_packet:
-            self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, chan_type, True)
+            self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, src_format, chan_type, True)
             self.prog.instr_dag.add_signal(self.rank, self, dst_chunkref, -1, ChannelType.none)
             self.prog.instr_dag.add_wait(dst, dst_chunkref, self, -1, ChannelType.none)
         else:
-            self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, chan_type)
+            self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, src_format, chan_type)
         return dst_chunkref
 
     def put(self, dst, buffer=None, index=-1, sendtb=-1, chan_type=ChannelType.memory):
-        return self._put(dst, buffer, index, sendtb, chan_type)
+        return self._put(dst, buffer, index, sendtb, DataFormat.raw, chan_type)
 
     def put_packet(
         self,
         dst,
         buffer=None,
         index=-1,
         sendtb=-1,
+        src_format=DataFormat.raw,
         chan_type=ChannelType.memory,
         temp_buffer=None,
         temp_buffer_index=-1,
     ):
         chunk_ref = self
-        if chan_type == ChannelType.port:
+        if chan_type == ChannelType.port and src_format == DataFormat.raw:
             assert temp_buffer is not None, "Need to specify a temporary buffer for port channels"
             chunk_ref = self._copy(
                 self.rank, temp_buffer, temp_buffer_index, sendtb, trans_from_packet=False, trans_to_packet=True
             )
-        return chunk_ref._put(dst, buffer, index, sendtb, chan_type, True)
+        return chunk_ref._put(dst, buffer, index, sendtb, src_format, chan_type, True)
 
     def get(self, src, buffer=None, index=-1, recvtb=-1, chan_type=ChannelType.memory):
         self.prog.check_buffer_exists(src, buffer)

python/mscclpp/language/types.py

@@ -86,6 +86,7 @@ class Instruction(Enum):
     reduce_send_packet = "rspkt"
     reduce_packet = "rpkt"
     put = "put"
+    read_put_packet = "rppkt"
     put_packet = "ppkt"
     put_with_signal = "pws"
     put_with_signal_and_flush = "pwsf"
@@ -175,3 +176,11 @@ def __hash__(self):
 
     def __repr__(self):
         return f"Op({self.inst}, {self.rank}, {self.src}, {self.dst}, step:{self.step}, tb:{self.tb})"
+
+
+class DataFormat(Enum):
+    raw = "raw"
+    packet = "packet"
+
+    def __str__(self):
+        return self.value

src/executor/execution_plan.cc

@@ -46,6 +46,8 @@ auto getOpType = [](const std::string& str) {
     return mscclpp::OperationType::READ_REDUCE_COPY_SEND;
   } else if (str == "ppkt") {
     return mscclpp::OperationType::PUT_PACKET;
+  } else if (str == "rppkt") {
+    return mscclpp::OperationType::READ_PUT_PACKET;
   } else if (str == "rspkt") {
     return mscclpp::OperationType::REDUCE_SEND_PACKET;
   } else if (str == "cpkt") {

src/include/execution_common.hpp

@@ -34,6 +34,7 @@ enum class OperationType : uint8_t {
   BARRIER,
   PUT,
   PUT_PACKET,
+  READ_PUT_PACKET,
   PUT_WITH_SIGNAL,
   PUT_WITH_SIGNAL_AND_FLUSH,
   GET,

src/include/execution_kernel.hpp

@@ -328,12 +328,53 @@ MSCCLPP_DEVICE_INLINE void handlePutPacket(size_t scratchSize, DeviceHandle<Memo
       return;
     }
     // For port channel, we assume src and dst are in packet format
+    // TODO: support non-packet format and remove packet format(packet format should be handle in handleReadPutPacket)
     uint32_t dstOffset = (dstOffsets[tid] << 1) + scratchBaseOffset;
     uint32_t srcOffset = (srcOffsets[tid] << 1) + scratchBaseOffset;
     portChannels[dstChannelIndexes[tid]].put(dstOffset, srcOffset, size << 1);
   }
 }
 
+template <typename PacketType>
+MSCCLPP_DEVICE_INLINE void handleReadPutPacket(int rank, void* scratch, size_t scratchSize,
+                                               DeviceHandle<SmChannel>* smChannels,
+                                               DeviceHandle<ProxyChannel>* proxyChannels, uint8_t* dstChannelIndexes,
+                                               uint32_t* dstOffsets, uint32_t* srcOffsets, int nDstChannels,
+                                               uint32_t size, ChannelType chType, uint32_t flag) {
+  const size_t scratchBaseOffset = flag & 0x1 ? 0 : scratchSize >> 1;
+  if (chType == ChannelType::MEMORY) {
+    size_t nPackets = size * 2 / sizeof(PacketType);
+    for (size_t pkt_idx = threadIdx.x; pkt_idx < nPackets; pkt_idx += blockDim.x) {
+      for (int ch_idx = 0; ch_idx < nDstChannels; ++ch_idx) {
+        PacketType* pkts = (PacketType*)((char*)scratch + scratchBaseOffset + srcOffsets[ch_idx] * 2);
+        PacketPayload<PacketType> data = pkts[pkt_idx].read(flag);
+        PacketType pkt(data, flag);
+        size_t offset = (scratchBaseOffset + dstOffsets[ch_idx] * 2) / sizeof(PacketType);
+        smChannels[dstChannelIndexes[ch_idx]].write(offset + pkt_idx, pkt);
+      }
+    }
+  } else if (chType == ChannelType::PORT) {
+    // Ensuring Data Is Ready
+    size_t nPackets = size * 2 / sizeof(PacketType);
+    for (size_t pkt_idx = threadIdx.x; pkt_idx < nPackets; pkt_idx += blockDim.x) {
+      for (int ch_idx = 0; ch_idx < nDstChannels; ++ch_idx) {
+        PacketType* pkts = (PacketType*)((char*)scratch + scratchBaseOffset + srcOffsets[ch_idx] * 2);
+        PacketPayload<PacketType> data = pkts[pkt_idx].read(flag);
+      }
+    }
+    __syncthreads();
+
+    // Putting the data
+    int ch_idx = threadIdx.x;
+    if (ch_idx >= nDstChannels) {
+      return;
+    }
+    uint32_t dstOffset = scratchBaseOffset + dstOffsets[ch_idx] * 2;
+    uint32_t srcOffset = scratchBaseOffset + srcOffsets[ch_idx] * 2;
+    proxyChannels[dstChannelIndexes[ch_idx]].put(dstOffset, srcOffset, size * 2);
+  }
+}
+
 template <typename T, typename PacketType, bool SendToRemote = true>
 MSCCLPP_DEVICE_INLINE void handleReduceSendPacket(T* dst, uint32_t dstOffsetByBytes, T* src, uint32_t srcOffsetByBytes,
                                                   T* inputBuff, size_t inputBuffSize, uint32_t* inputOffsets, int nSrcs,
@@ -571,6 +612,9 @@ __global__ void executionKernel([[maybe_unused]] int rank /*for debug*/, T* inpu
       handleReadReduceCopySend(dst, op.dstOffset, src, op.srcOffset, memoryChannels, op.outputChannelIndexes,
                                op.inputChannelIndexes, op.outputOffsets, op.inputOffsets, op.nOutputs, op.nInputs,
                                op.size, false);
+    } else if (op.type == OperationType::READ_PUT_PACKET) {
+      handleReadPutPacket<PacketType>(rank, scratch, scratchSize, memoryChannels, portChannels, op.outputChannelIndexes,
+                                      op.outputOffsets, op.inputOffsets, op.nOutputs, op.size, op.channelType, flag);
     } else if (op.type == OperationType::PUT_PACKET) {
       handlePutPacket<PacketType>(scratchSize, memoryChannels, portChannels, op.outputChannelIndexes, op.outputOffsets,
                                   op.inputOffsets, op.nOutputs, op.size, op.channelType, flag);