From a3421265431b429df87318d1a5c6cc796c41c265 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 20 May 2026 23:24:33 +0200 Subject: [PATCH] polish(crt-train): branchless boundary loads + remove misleading factor==0 comment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Code-quality reviewer flagged two non-blocking polish items in the output_smoothness kernel's backward pass: - I1: thread-divergent if-gated loads contradicted the file's stated branchless-design intent. Fixed via clamp-then-load: at k=0 / k=K-1 read self for the spurious neighbour; lm/rm multipliers zero the contribution. Every lane executes both loads; no divergence on boundaries. - I2: the "Skip work when factor == 0" comment described a non-feature (no such branch existed). Removed. Behaviour-equivalent change — same loss and same gradient at every position; the change only affects warp-execution patterns at the K=0 and K=K-1 strips. --- crates/ml-alpha/cuda/output_smoothness.cu | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/ml-alpha/cuda/output_smoothness.cu b/crates/ml-alpha/cuda/output_smoothness.cu index b91274b46..1e784e401 100644 --- a/crates/ml-alpha/cuda/output_smoothness.cu +++ b/crates/ml-alpha/cuda/output_smoothness.cu @@ -146,18 +146,18 @@ extern "C" __global__ void output_smoothness_loss_and_grad( const int k = i / stride_k; const float lam = s_lambda[h]; const float factor = 2.0f * lam * inv_n_pairs; - // Skip work when factor == 0 (e.g. lambda all zero); still must - // not branch on host — the multiply by zero is the no-op. const float p_j = probs[i]; - float p_jm1 = 0.0f; - float p_jp1 = 0.0f; - // Read p[j-1] only if k >= 1; otherwise no contribution from - // the left pair (we encode this by treating p_jm1 == p_j so the - // (p[j] - p[j-1]) term in Δ contributes 0). Branchless via select. const bool has_left = (k >= 1); const bool has_right = (k <= K - 2); - if (has_left) p_jm1 = probs[i - stride_k]; - if (has_right) p_jp1 = probs[i + stride_k]; + // Branchless boundary: at k=0 we read self for the "left" slot + // and at k=K-1 we read self for the "right" slot; both indices + // stay in-bounds and the lm/rm multipliers below zero the + // spurious contribution. Every lane executes both loads — no + // thread divergence on boundaries. + const int idx_left = has_left ? (i - stride_k) : i; + const int idx_right = has_right ? (i + stride_k) : i; + const float p_jm1 = probs[idx_left]; + const float p_jp1 = probs[idx_right]; // Δ(j, b, h): // left contribution = (p[j] - p[j-1]) if has_left else 0 // right contribution = -(p[j+1] - p[j]) if has_right else 0