Change MSCCL kernel signature to allow kernel arguments be preloaded via SGPR (#911)

* Adding a script that will download/compile/run TransferBench/RCCL/UCX/RCCL-tests/RCCL-Unittests/hip-mpi-testsuite (#895)

Co-authored-by: Pedram Alizadeh <pmohamma@banff-pla-r27-05.pla.dcgpu>

* Only build gfx941

* demo

* fine tune malloc

* Fix merge errors

* Fix merge errors

* Disable parallel build

* Adopt --amdgpu-kernarg-preload-count

* Revert "Adding a script that will download/compile/run TransferBench/RCCL/UCX/RCCL-tests/RCCL-Unittests/hip-mpi-testsuite (#895)"

This reverts commit f5e252dddf02a41b4d1bc512f306f45f97166304.

* Revert CMake changes.

* NPKIT changes.

* Remove some license declarations.

* Address code review feedbacks on msccl_kernel_impl.h

* Update CMakeLists.txt

* Add CMake logic to check the existence of --amdgpu-kernarg-preload-count

* Fix NPKIT trace logic.

---------

Co-authored-by: Pedram Alizadeh <pmohamma@amd.com>
Co-authored-by: Pedram Alizadeh <pmohamma@banff-pla-r27-05.pla.dcgpu>
Co-authored-by: Ziyue Yang <ziyyang@microsoft.com>

[ROCm/rccl commit: 7ee5c1c28b]

projects/rccl/src/include/msccl/msccl_kernel.h

@@ -9,7 +9,7 @@
 #define MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto) mscclKernel_##devredop##_##type##_##proto
 
 #define MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE_PROTO(devredop, type, proto) \
-__global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork work);
+__global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work);
 
 #define MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, type) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE_PROTO(devredop, type, LL) \

projects/rccl/src/include/msccl/msccl_struct.h

@@ -36,6 +36,8 @@
 #define MSCCL_LOCAL_COPY 6
 #define MSCCL_REDUCE 7
 
+#define MSCCL_WORK_FIFO_DEPTH (64 << 10)
+
 struct mscclTransmission {
   int16_t dependencePointer; // index to the first dependence
   int16_t numDependencies; // dependencePointer+numDependencies indicate the last dependence
@@ -210,21 +212,36 @@ struct mscclStatus {
   bool graphEnabled;
   bool graphFirstKernel;
   bool needsProxy;
+  uint64_t workFifoDepth;
+  struct mscclWork* workFifo;
+  uint32_t* workFifoDone;
+  uint32_t workFifoSent;
+  uint32_t workFifoSentPerChannel[MAXCHANNELS];
+  uint32_t workFifoAckdMin;
 };
 
-struct alignas(16) mscclWork {
+#pragma pack(push)
+#pragma pack(8)
+
+struct mscclWork {
   volatile struct mscclFlag *syncFlags;
   void *scratchBuffer;
   const void *sendBuff;
   void *recvBuff;
-  size_t count;
+  uint32_t* workFifoDone;
+  size_t sizePerMscclChunk;
   uint64_t redOpArg;
   uint32_t workIndex;
-  int nChunksPerLoop;
   uint32_t maxAllowedCount;
+  uint32_t workFifoDoneAck;
+  int nChunksPerLoop;
   bool hasReduce;
   bool redOpArgIsPtr;
+  uint32_t pad[1];
 };
+static_assert(sizeof(struct mscclWork) % 16 == 0, "mscclWork needs to be 16B aligned");
+
+#pragma pack(pop)
 
 struct mscclShmemData {
   struct mscclThreadBlock mscclTB;