[mlir-tensorrt] NFC: [executor] update NCCL debug logging in runtime

Adds some additional debugging logging in the NCCL executor module.

GitOrigin-RevId: b628f4613e87e148ed444342932dc1cf909851e3

Author: christopherbate

mlir-tensorrt/executor/lib/Runtime/Backend/Lua/Modules/CUDA/CUDAModule.cpp

@@ -227,6 +227,7 @@ registerCudaMemoryManagementOps(sol::state_view &lua,
   };
 
   lua["__cuda_stream_sync"] = [](sol::this_state state, CudaStreamPtr stream) {
+    MTRT_DBG("__cuda_stream_sync @ {0}", reinterpret_cast<void *>(stream.ptr));
     ADD_CUDA_MODULE_RANGE("cuda_stream_sync");
     SET_LUA_ERROR_IF_CUDART_ERROR(cudaStreamSynchronize(stream), state);
   };
@@ -439,7 +440,8 @@ registerCudaMemoryManagementOps(sol::state_view &lua,
                      size_t srcOffset, uintptr_t dest, size_t destOffset,
                      size_t numBytes) {
         ADD_CUDA_MODULE_RANGE("cuda_memcpy_host_pinned2device");
-        MTRT_DBGF("cuda_memcpy_h2d %lu bytes from 0x%lx + %lu to 0x%lx + %lu",
+        MTRT_DBGF("__cuda_memcpy_host_pinned2device: %lu bytes from 0x%lx + "
+                  "%lu to 0x%lx + %lu",
                   numBytes, src, srcOffset, dest, destOffset);
         void *srcPtr = reinterpret_cast<void *>(src + srcOffset);
         void *dstPtr = reinterpret_cast<void *>(dest + destOffset);
@@ -475,7 +477,9 @@ registerCudaMemoryManagementOps(sol::state_view &lua,
                  "expected src to be a device ptr and dest to be a host ptr");
         }
 #endif
-        MTRT_DBGF("executor_memcpy device-host %lu bytes", numBytes);
+        MTRT_DBGF("__cuda_memcpy_device2host_pinned: %lu bytes from 0x%lx + "
+                  "%lu to 0x%lx + %lu",
+                  numBytes, src, srcOffset, dest, destOffset);
         SET_LUA_ERROR_IF_CUDART_ERROR(cudaMemcpyAsync(dstPtr, srcPtr, numBytes,
                                                       cudaMemcpyDeviceToHost,
                                                       stream),

mlir-tensorrt/executor/lib/Runtime/Backend/Lua/Modules/NCCL/NCCLModule.cpp

@@ -268,6 +268,10 @@ static void registerNcclOps(sol::state_view &lua, ResourceTracker *tracker) {
   lua["__nccl_all_reduce_" #opsuffix "_" #typesuffix] =                        \
       [](sol::this_state state, ExecPtr sendbuff, ExecPtr recvbuff,            \
          size_t count, uintptr_t communicator, CudaStreamPtr stream) {         \
+        MTRT_DBG("__nccl_all_reduce_" #opsuffix "_" #typesuffix                \
+                 ": count={0} send={1} recv={2}",                              \
+                 count, reinterpret_cast<void *>(sendbuff),                    \
+                 reinterpret_cast<void *>(recvbuff));                          \
         auto comm = reinterpret_cast<NcclCommunicator *>(communicator);        \
         SET_LUA_ERROR_IF_NCCL_ERROR(                                           \
             ncclAllReduce(reinterpret_cast<void *>(sendbuff),                  \
@@ -285,6 +289,10 @@ static void registerNcclOps(sol::state_view &lua, ResourceTracker *tracker) {
   lua["__nccl_reduce_scatter_" #opsuffix "_" #typesuffix] =                    \
       [](sol::this_state state, ExecPtr sendbuff, ExecPtr recvbuff,            \
          size_t recvcount, uintptr_t communicator, CudaStreamPtr stream) {     \
+        MTRT_DBG("__nccl_reduce_scatter_" #opsuffix "_" #typesuffix            \
+                 ": count={0} sendbuff={1} recvbuff={2}",                      \
+                 recvcount, reinterpret_cast<void *>(sendbuff),                \
+                 reinterpret_cast<void *>(recvbuff));                          \
         auto *comm = reinterpret_cast<NcclCommunicator *>(communicator);       \
         SET_LUA_ERROR_IF_NCCL_ERROR(                                           \
             ncclReduceScatter(reinterpret_cast<void *>(sendbuff),              \
@@ -338,6 +346,13 @@ static void registerNcclOps(sol::state_view &lua, ResourceTracker *tracker) {
                              size_t numBytes, uintptr_t communicator,
                              CudaStreamPtr stream) {
     auto *comm = reinterpret_cast<NcclCommunicator *>(communicator);
+    MTRT_DBG("__nccl_permute[{6}/{7}]: send {0} bytes @ {1} to {2}, recv {0} "
+             "bytes @ "
+             "{3} from {4}, comm @{5}",
+             numBytes, reinterpret_cast<void *>(sendbuff), sendId,
+             reinterpret_cast<void *>(recvbuff), recvId,
+             reinterpret_cast<void *>(comm->comm), comm->rank, comm->numRanks);
+
     if (recvId == -1) {
       // Zero out recvbuff if not receiving.
       SET_LUA_ERROR_IF_CUDA_ERROR(