Support fused all reduce and elementwise operations (#1729)

* Support fused all reduce and elementwise operations

Add additional "acc" parameter to RCCL Replayer logs

Add flag which indicates availability of new API

* Fix Recorder json parsing

* Remove unreachable code

* Remove extra acc pointer check

* .

* Revert "[DEVICE] Adding ability to choose unroll factor at runtime (#1734)"

This reverts commit 4cadf3597c.

* Use noinline to reduce kernels linking time

* Don't use noinline for gfx942 and gfx950 to avoid perf regression

---------

Co-authored-by: AtlantaPepsi <timhu102@amd.com>
Co-authored-by: BertanDogancay <bertan.dogancay@gmail.com>

[ROCm/rccl commit: 9a4213356d]

projects/rccl/src/include/api_trace.h

@@ -31,7 +31,7 @@
 #define RCCL_API_TRACE_VERSION_MAJOR 0
 
 // should be increased every time new members are added to existing dispatch tables
-#define RCCL_API_TRACE_VERSION_PATCH 0
+#define RCCL_API_TRACE_VERSION_PATCH 1
 
 #if !defined(RCCL_EXTERN_C_INIT)
 #    ifdef __cplusplus
@@ -61,6 +61,10 @@ typedef ncclResult_t (*ncclAllReduce_fn_t)(const void* sendbuff, void* recvbuff,
                                            size_t count, ncclDataType_t datatype,
                                            ncclRedOp_t op, struct ncclComm* comm,
                                            hipStream_t stream);
+typedef ncclResult_t (*ncclAllReduceWithBias_fn_t)(const void* sendbuff, void* recvbuff,
+                                           size_t count, ncclDataType_t datatype,
+                                           ncclRedOp_t op, struct ncclComm* comm,
+                                           hipStream_t stream, const void* acc);
 typedef ncclResult_t (*ncclAllToAll_fn_t)(const void* sendbuff, void* recvbuff,
                                           size_t count, ncclDataType_t datatype,
                                           ncclComm_t comm, hipStream_t stream);
@@ -194,6 +198,7 @@ typedef struct rcclApiFuncTable
     mscclUnloadAlgo_fn_t          mscclUnloadAlgo_fn;
     ncclCommRegister_fn_t         ncclCommRegister_fn;
     ncclCommDeregister_fn_t       ncclCommDeregister_fn;
+    ncclAllReduceWithBias_fn_t    ncclAllReduceWithBias_fn;
 
 } rcclApiFuncTable;
 

projects/rccl/src/include/comm.h

@@ -195,6 +195,7 @@ struct ncclTaskColl {
   ncclFunc_t func;
   void const* sendbuff;
   void* recvbuff;
+  void const* acc;
   size_t count;
   int root;
   ncclDataType_t datatype;

projects/rccl/src/include/device.h

@@ -310,6 +310,7 @@ struct alignas(16) ncclDevWorkColl {
   uint16_t pivotA2ANumBiRings:15, profilerEnabled:1;
   void* recvbuff;
   void* sendbuff;
+  void *acc;
   uintptr_t sendbuffOffset;
   uintptr_t recvbuffOffset;
   uintptr_t* sendbuffRmtAddrs;

projects/rccl/src/include/info.h

@@ -29,6 +29,7 @@ struct ncclInfo {
   // Algorithm details
   int chunkSteps;
   int sliceSteps;
+  const void* acc;
 };
 
 #endif

projects/rccl/src/include/nccl_common.h

@@ -74,5 +74,10 @@ typedef enum {
 
 #define NCCL_ALGO_PROTO_IGNORE -1.0
 
+#define NCCL_NUM_UNROLLS 3 // 1/2/4
+#define NCCL_UNROLL_1 0
+#define NCCL_UNROLL_2 1
+#define NCCL_UNROLL_4 2
+
 #define NCCL_NUM_FLOATS 6 // half/float/double/rccl_bfloat16/rccl_float8/rccl_bfloat8
 #endif

projects/rccl/src/include/recorder.h

@@ -16,6 +16,7 @@ typedef enum {
   rrAllGather,
   rrReduceScatter,
   rrAllReduce,
+  rrAllReduceWithBias,
   rrSend,
   rrRecv,
   rrAllToAll,
@@ -51,6 +52,7 @@ constexpr const char* rcclCallStr[]
   "AllGather",
   "ReduceScatter",
   "AllReduce",
+  "AllReduceWithBias",
   "Send",
   "Recv",
   "AllToAll",
@@ -94,6 +96,7 @@ struct rcclApiCall {
   uint64_t              opCount = 0;
   const void*           sendbuff = NULL;
   void*                 recvbuff = NULL;
+  const void*           acc = NULL;
   void*                 sendPtrBase = NULL;
   void*                 recvPtrBase = NULL;
   size_t                sendPtrExtent = 0;