Improve collective trace

src/include/collectives.h

@@ -9,6 +9,9 @@
 #define NCCL_COLLECTIVES_H_
 
 #define FUNC_INDEX_P2P (NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps)
+#define FUNC_INDEX_A2A (FUNC_INDEX_P2P+1)
+#define FUNC_INDEX_A2AV (FUNC_INDEX_P2P+2)
+
 #define FUNC_INDEX(func, redop, ncclType, al, pr) ((func >= NCCL_NUM_FUNCTIONS) \
   ? (func-NCCL_NUM_FUNCTIONS+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps) \
   : ((((((func)*ncclNumOps + (redop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr)))

src/include/devcomm.h

@@ -286,7 +286,18 @@ struct ncclCollTrace {
   uint32_t data_0;
   uint64_t timeStamp;
   uint64_t opCount;
-  uint64_t data_1;
+  union {
+    uint64_t data_1;
+    struct {
+      uint16_t nThreads;
+      uint8_t bid;
+      uint8_t nChannels;
+    } coll;
+    struct {
+      uint16_t nThreads;
+      uint16_t delta;
+    } p2p;
+  };
 };
 static_assert(sizeof(struct ncclCollTrace) == 8*sizeof(int), "ncclCollTrace must have a pow2 size");