uvm: coalesce TLB batch ranges and pick least-busy channels

kernel-open/nvidia-uvm/uvm_channel.c

@@ -401,6 +401,43 @@ static bool test_claim_and_lock_channel(uvm_channel_t *channel,
     return false;
 }
 
+// Returns a channel that can satisfy the reserve, or NULL. Does not claim.
+static uvm_channel_t *channel_pick_least_busy_locked(uvm_channel_pool_t *pool,
+                                                     NvU32 num_gpfifo_entries,
+                                                     uvm_channel_reserve_type_t reserve_type,
+                                                     bool skip_locked_for_push)
+{
+    NvU32 n;
+    NvU32 best_available = 0;
+    uvm_channel_t *best = NULL;
+
+    uvm_channel_pool_assert_locked(pool);
+    UVM_ASSERT(!skip_locked_for_push || g_uvm_global.conf_computing_enabled);
+
+    for (n = 0; n < pool->num_channels; n++) {
+        NvU32 index = (pool->next_hint + n) % pool->num_channels;
+        uvm_channel_t *channel = &pool->channels[index];
+        NvU32 available;
+
+        if (skip_locked_for_push && uvm_channel_is_locked_for_push(channel))
+            continue;
+
+        if (reserve_type == UVM_CHANNEL_RESERVE_WITH_P2P && channel->suspended_p2p)
+            continue;
+
+        available = channel_get_available_gpfifo_entries(channel);
+        if (available < num_gpfifo_entries)
+            continue;
+
+        if (best == NULL || available > best_available) {
+            best = channel;
+            best_available = available;
+        }
+    }
+
+    return best;
+}
+
 // Reserve, or release, all channels in the given pool.
 //
 // One scenario where reservation of the entire pool is useful is key rotation,
@@ -512,7 +549,6 @@ static NV_STATUS channel_reserve_and_lock_in_pool(uvm_channel_pool_t *pool,
 {
     uvm_channel_t *channel;
     uvm_spin_loop_t spin;
-    NvU32 index;
     NV_STATUS status;
 
     UVM_ASSERT(pool);
@@ -529,20 +565,15 @@ static NV_STATUS channel_reserve_and_lock_in_pool(uvm_channel_pool_t *pool,
     // uvm_channel_end_push() routine.
     uvm_down(&pool->conf_computing.push_sem);
 
-    // At least one channel is unlocked. We check if any unlocked channel is
-    // available, i.e., if it has free GPFIFO entries.
-
+    // At least one channel is unlocked. Check if any unlocked channel has a
+    // free GPFIFO entry.
     channel_pool_lock(pool);
-
-    for_each_clear_bit(index, pool->conf_computing.push_locks, pool->num_channels) {
-        channel = &pool->channels[index];
-
-        if (try_claim_channel_locked(channel, 1, reserve_type)) {
-            lock_channel_for_push(channel);
-            goto done;
-        }
+    channel = channel_pick_least_busy_locked(pool, 1, reserve_type, true);
+    if (channel != NULL && try_claim_channel_locked(channel, 1, reserve_type)) {
+        lock_channel_for_push(channel);
+        pool->next_hint = (uvm_channel_index_in_pool(channel) + 1) % pool->num_channels;
+        goto done;
     }
-
     channel_pool_unlock(pool);
 
     // No channels are available. Update and check errors on all channels until
@@ -593,13 +624,15 @@ static NV_STATUS channel_reserve_in_pool(uvm_channel_pool_t *pool,
     if (g_uvm_global.conf_computing_enabled)
         return channel_reserve_and_lock_in_pool(pool, reserve_type, channel_out);
 
-    uvm_for_each_channel_in_pool(channel, pool) {
-        // TODO: Bug 1764953: Prefer idle/less busy channels
-        if (try_claim_channel(channel, 1, reserve_type)) {
-            *channel_out = channel;
-            return NV_OK;
-        }
+    channel_pool_lock(pool);
+    channel = channel_pick_least_busy_locked(pool, 1, reserve_type, false);
+    if (channel != NULL && try_claim_channel_locked(channel, 1, reserve_type)) {
+        pool->next_hint = (uvm_channel_index_in_pool(channel) + 1) % pool->num_channels;
+        channel_pool_unlock(pool);
+        *channel_out = channel;
+        return NV_OK;
     }
+    channel_pool_unlock(pool);
 
     uvm_spin_loop_init(&spin);
     while (1) {

kernel-open/nvidia-uvm/uvm_channel.h

@@ -211,6 +211,10 @@ typedef struct
     // Number of elements in the channel array
     NvU32 num_channels;
 
+    // Next channel to consider when selecting a least-busy channel.
+    // Protected by the pool lock.
+    NvU32 next_hint;
+
     // Index of the engine associated with the pool (index is an offset from the
     // first engine of the same engine type.)
     unsigned engine_index;

kernel-open/nvidia-uvm/uvm_channel_test.c

@@ -977,6 +977,105 @@ static NV_STATUS test_conf_computing_channel_selection(uvm_va_space_t *va_space)
     return status;
 }
 
+// Concurrent uvm_push_begin calls on an idle multi-channel pool should reserve
+// distinct channels.
+static NV_STATUS test_least_busy_channel_reserve(uvm_va_space_t *va_space)
+{
+    NV_STATUS status = NV_OK;
+    uvm_push_t *pushes = NULL;
+    uvm_gpu_t *gpu = NULL;
+    NvU32 num_begun = 0;
+    NvU32 num_ended = 0;
+
+    uvm_thread_context_lock_disable_tracking();
+
+    for_each_va_space_gpu(gpu, va_space) {
+        uvm_channel_type_t channel_type;
+
+        // Nested begins are illegal with key rotation. Disable it for the
+        // duration of this GPU's concurrent pushes; a no-op when CC is off.
+        uvm_conf_computing_disable_key_rotation(gpu);
+
+        for (channel_type = 0; channel_type < UVM_CHANNEL_TYPE_COUNT; channel_type++) {
+            NvU32 i;
+            NvU32 num_pushes;
+            uvm_channel_t *channel;
+            uvm_channel_pool_t *pool = gpu->channel_manager->pool_to_use.default_for_type[channel_type];
+
+            // SEC2/WLC/LCIC defaults are NULL when Confidential Computing is off.
+            if (pool == NULL)
+                continue;
+
+            // Skip LCIC channels as those can't accept any pushes
+            if (uvm_channel_pool_is_lcic(pool))
+                continue;
+
+            if (pool->num_channels < 2)
+                continue;
+
+            num_pushes = pool->num_channels;
+
+            // Reclaim completed GPFIFO so available counts match an idle pool.
+            uvm_for_each_channel_in_pool(channel, pool)
+                TEST_NV_CHECK_GOTO(uvm_channel_wait(channel), error);
+
+            pushes = uvm_kvmalloc_zero(sizeof(*pushes) * num_pushes);
+            TEST_CHECK_GOTO(pushes != NULL, error);
+
+            num_begun = 0;
+            num_ended = 0;
+
+            for (i = 0; i < num_pushes; i++) {
+                uvm_push_t *push = &pushes[i];
+                status = uvm_push_begin(gpu->channel_manager, channel_type, push, "least-busy push %u", i);
+                TEST_NV_CHECK_GOTO(status, error);
+                num_begun++;
+            }
+
+            for (i = 0; i < num_pushes; i++) {
+                NvU32 j;
+
+                for (j = 0; j < i; j++)
+                    TEST_CHECK_GOTO(pushes[i].channel != pushes[j].channel, error);
+            }
+
+            for (i = 0; i < num_pushes; i++) {
+                uvm_push_t *push = &pushes[i];
+                status = uvm_push_end_and_wait(push);
+                num_ended++;
+                TEST_NV_CHECK_GOTO(status, error);
+            }
+
+            uvm_kvfree(pushes);
+            pushes = NULL;
+            num_begun = 0;
+            num_ended = 0;
+        }
+
+        uvm_conf_computing_enable_key_rotation(gpu);
+    }
+
+    uvm_thread_context_lock_enable_tracking();
+
+    return status;
+
+error:
+    if (pushes != NULL) {
+        NvU32 i;
+
+        for (i = num_ended; i < num_begun; i++)
+            uvm_push_end(&pushes[i]);
+    }
+
+    if (gpu != NULL)
+        uvm_conf_computing_enable_key_rotation(gpu);
+
+    uvm_thread_context_lock_enable_tracking();
+    uvm_kvfree(pushes);
+
+    return status;
+}
+
 static NV_STATUS test_channel_iv_rotation(uvm_va_space_t *va_space)
 {
     uvm_gpu_t *gpu;
@@ -1640,6 +1739,10 @@ NV_STATUS uvm_test_channel_sanity(UVM_TEST_CHANNEL_SANITY_PARAMS *params, struct
     if (status != NV_OK)
         goto done;
 
+    status = test_least_busy_channel_reserve(va_space);
+    if (status != NV_OK)
+        goto done;
+
     status = test_channel_iv_rotation(va_space);
     if (status != NV_OK)
         goto done;