TL/CUDA: free bar in progress

openucx · Nov 21, 2024 · ce8d0c2 · ce8d0c2
1 parent 109a5dc
commit ce8d0c2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/src/components/tl/cuda/bcast/bcast_linear.c b/src/components/tl/cuda/bcast/bcast_linear.c
@@ -82,7 +82,7 @@ static inline ucc_status_t root_find_free_barrier(ucc_tl_cuda_task_t *task)
         // try to set user specified tag to mark that this barrier is used by this task
         if (ucc_atomic_cswap64(&curr_bar->tag, UCC_TAG_FREE,
                                task->bcast_linear.key) == UCC_TAG_FREE) {
-            ucc_print("found free barrier: %d marked with tag: %ld", i,
+            ucc_print("found free barrier: %p idx: %d marked with tag: %ld", curr_bar, i,
                       curr_bar->tag);
             // free
             task->bar = curr_bar;
@@ -126,7 +126,6 @@ static inline ucc_status_t peer_find_free_barrier(ucc_tl_cuda_task_t *task)
             }
             found                    = true;
             task->coll_id            = i + max_concurrent;
-            // task->bcast_linear.stage = STAGE_SYNC;
             break;
         }
     }
@@ -292,6 +291,9 @@ void ucc_tl_cuda_bcast_linear_progress(ucc_coll_task_t *coll_task)
                 // finish
                 ucc_tl_cuda_put_sync_root(task, task->bcast_linear.root);
                 task->super.status = UCC_OK;
+                // set barrier free to unlock others, this is roots responsibility
+                ucc_print("Free bar: %p with tag: %ld", task->bar, task->bar->tag);
+                task->bar->tag = UCC_TAG_FREE;                
                 break;
             }
         default:

diff --git a/src/components/tl/cuda/tl_cuda_coll.h b/src/components/tl/cuda/tl_cuda_coll.h
@@ -86,11 +86,6 @@ static inline ucc_tl_cuda_task_t *ucc_tl_cuda_task_get(ucc_tl_cuda_team_t *team)
 static inline void ucc_tl_cuda_task_put(ucc_tl_cuda_task_t *task)
 {
     UCC_TL_CUDA_PROFILE_REQUEST_FREE(task);
-
-    if (UCC_TL_TEAM_RANK(TASK_TEAM(task)) == task->bcast_linear.root) {
-        task->bar->tag = UCC_TAG_FREE;
-    }
-
     ucc_mpool_put(task);
 }