ARCH/X86: Rework AVX2 NT buffer-transfer kernels and dispatch

[arun-chandran-edarath]

What?

Rework the AVX2 NT buffer-transfer kernels to handle edges and tails internally, simplify dispatch, update AMD thresholds, and select ERMS for eligible NT source copies. Extend coverage for both vectorized and ERMS paths.

Why?

Feature: splitting as per comments from #11696

How?

#11696 is broken to logical PRs, this is the second in the series

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[arun-chandran-edarath]

@roiedanino @tvegas1 please check.

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[svc-nvidia-pr-review]

NT_BUFFER_TRANSFER_MIN also zeroes builtin_memcpy_max, which disables ERMS in ucs_memcpy_relaxed() for all copy sizes, not just NT transfers — including on Intel where NT is off. A user tuning only the NT threshold loses ERMS for regular copies. Can we scope this so it only affects the NT dispatch path, or gate it on the NT vendor/threshold being active?

[arun-chandran-edarath]

Correct, an explicit NT_BUFFER_TRANSFER_MIN intentionally closes the outer ERMS window for all sizes. Transfers below the configured NT threshold therefore use memcpy(), while eligible transfers reach ucs_x86_nt_buffer_transfer(), where ERMS remains available internally for NT_SOURCE.

This is needed because outer ERMS selects on fragment len before NT eligibility is evaluated on total_len. Retaining or clamping the outer window would allow an eligible transfer’s fragments to return through ERMS before its explicit NT policy and copy hint are considered. On Intel, the automatic NT threshold remains disabled; setting it explicitly opts into this behavior.

[svc-nvidia-pr-review]

src/ucs/arch/x86_64/cpu_nt_avx2.inl

The rewrite removes the aocl-libmem attribution comment (adaptation of the memcpy code from https://github.com/amd/aocl-libmem). Is the new implementation a clean rewrite that no longer derives from that source?

[arun-chandran-edarath]

src/ucs/arch/x86_64/cpu_nt_avx2.inl

The rewrite removes the aocl-libmem attribution comment (adaptation of the memcpy code from https://github.com/amd/aocl-libmem). Is the new implementation a clean rewrite that no longer derives from that source?

Yes, this is intentional. The attribution applied to the previous ucs_x86_copy_bytes_le_128()-based implementation. Its leading-zero switch and special len <= 128 dispatch were removed.

The current AVX2 kernels were reworked using the AVX-512 kernel design from the original PR #11696 as their structural reference: a 64-byte-aligned prefix, explicit 256-byte body, 64-byte drain, and sub-cache-line tail. The AVX-512 implementation is not part of this PR; it is planned for the third PR in the split series.

The rewritten AVX2 code does not retain AOCL-LibMem’s control flow or vector-loop implementation. The remaining overlapping head/tail handling for a short tail is a standard bounded-copy technique

[svc-ucx]

🤖 CI Triage Agent — UCX PR (AddressSanitizer roce on worker 0) · commit e7e3c62d

TL;DR: This job did not fail on any code error — the Azure Pipelines agent was killed mid-compile ("The agent has received a shutdown signal"), so the ASAN roce job was cancelled while still building gtest objects. Re-run the pipeline; there is nothing to fix in the PR.

Full analysis

Summary: The "AddressSanitizer roce on worker 0" job was cancelled by the Azure DevOps agent host during the make phase of test/gtest, before any test executed.

Root cause: Infrastructure, not code. The log shows a clean, uninterrupted build: libucp.so linked successfully at 10:29:37, tools/perf/profile/examples/apps all built, gtest googletest lib linked at 10:32:26, and gtest object compilation proceeding normally. The last application output is at 10:36:11.4260362Z compiling uct/test_amo_add_xor.cc, immediately followed at 10:36:11.9398211Z by:

##[error]The agent has received a shutdown signal. This can happen when the agent
         service is stopped, or a manually started agent is canceled.
##[error]The Operation will be canceled.
##[error]The operation was canceled.

Checking for a hang: the largest gap between consecutive log lines is ~39 s (10:30:36 → 10:31:15, the uct_tests.cc ASAN compile), which is normal for -O3 -fsanitize=address C++ compilation. There is continuous output every few seconds right up to the kill, so the process was neither hung nor starved for time — and it was only ~7 minutes into the step, far from any plausible timeout. There are zero compiler errors, zero linker errors, and zero gtest output (the gtest binary was never even linked, let alone run), so no ASAN report or test failure exists in this build.

The signature is an agent-side termination: agent VM preemption/reclamation, agent service restart, or Azure auto-cancelling the run because a newer commit was pushed to PR #11895 superseding commit e7e3c62. Note this is distinct from a compiler OOM, which would surface as cc1plus: out of memory / internal compiler error / Killed from make, none of which appear.

Implicated commit: none — [REDACTED:Hex High Entropy String] is not implicated; the failure occurred before any of its code was exercised.

File: n/a — no source file is implicated (last compiled unit test/gtest/uct/test_amo_add_xor.cc completed the step normally and is unrelated to the cancellation).

Suggested fix: Re-queue the "UCX PR" pipeline for PR #11895 (or push an empty commit) to get a clean signal on the ASAN roce job. Do not raise any timeout — the step was killed at ~7 minutes with continuous progress, so a time-limit change would have no effect. If this recurs across multiple PRs on the same agent pool, escalate to CI infrastructure to check for agent preemption/host restarts on the worker 0 agent, and consider enabling cancelTimeoutInMinutes plus build-artifact upload of the partial log so superseded-vs-crashed runs can be distinguished.

Related: PR #11895 (this build). No matching issues found for the "agent has received a shutdown signal" signature in the openucx/ucx tracker.

🛡️ This comment had 1 potential secret(s) redacted (Hex High Entropy String). See request_id 5b30c438-9e15-4daf-826b-5194f6ad01ff in the triage console for the audit trail.

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[svc-nvidia-pr-review]

Residual coverage gap: no test exercises a fragmented transfer where total_len >= nt_buffer_transfer_min but len < 3072 (the new memcpy early return) through ucs_memcpy_relaxed(), which is where the builtin_memcpy_max = 0 interaction actually shows up; the new gtests call ucs_x86_nt_buffer_transfer() directly.

Residual coverage gap: the ERMS NT_SOURCE test only runs where ENABLE_BUILTIN_MEMCPY is on and builtin_memcpy_min is forced to 0; on AMD hardware the default builtin_memcpy_min is UCS_MEMUNITS_INF, so CI on AMD never covers the ERMS branch with production settings.

[svc-ucx]

🤖 CI Triage Agent — UCX PR (Tests gpu on worker 0) · commit 4f35d6e5

TL;DR: The test_fuzzy_match.py step failed because ucx_info emitted a CUDA driver error line (cuDeviceGetCount(...) failed: initialization error) as the first line of output, and the test only inspects output_lines[0] when looking for the "unused environment variable" warning. Fix by scanning all output lines for the warning (and separately investigate the broken CUDA driver state on swx-rdmz-ucx-gpu-02).

Full analysis

Summary: Azure job "Tests gpu on worker 0" (./contrib/test_jenkins.sh, task test_unused_env_var) exited 1 when test/apps/test_fuzzy_match.py raised "unused vars" message was not found.

Root cause: Two-part, neither related to PR #11895 (which is an AVX2 NT-copy dispatch rework and touches no CUDA/parser code):

1. On the GPU worker, the CUDA driver failed to initialize — the log shows cuda_ctx.c:52 UCX ERROR cuDeviceGetCount(&num_devices) failed: initialization error (CUDA_ERROR_NOT_INITIALIZED) printed by uct_cuda_ctx_primary_push_first_active() during ucx_info -u m -w worker creation. This is a node/container GPU-availability problem, not a code change in this PR.
2. test_fuzzy_match.py is brittle: get_fuzzy_matches() takes warn_msg = output_lines[0] and regex-matches only that line. The expected warning was present (line 2 of the output: parser.c:2578 UCX WARN unused environment variable: UCX_LOF_LEVEL (maybe: UCX_LOG_LEVEL?)), but because the CUDA ERROR line landed at index 0, the regex failed and the test aborted.

Implicated commit: unknown — the failing commit 4f35d6e (ntbt_dispatch_rework, PR #11895) is not implicated; no recent commits touch test/apps/test_fuzzy_match.py, and src/uct/cuda/base/cuda_ctx.c was last changed in 58be745 (Raul Akhmetshin, #11366) long before this build.

File: test/apps/test_fuzzy_match.py:67 (and the error source at src/uct/cuda/base/cuda_ctx.c:52)

Suggested fix:

• Make the test tolerant of preceding log lines, e.g. replace line 67 with a search over all lines:

  warn_msg = next((l for l in output_lines
                   if re.match('.*unused environment variables?: ', l)), '')

  (keep the existing if not warn_match: diagnostic path). Optionally also drop lines matching UCX ERROR|UCX DEBUG before parsing.
• Independently, triage the GPU node: cuDeviceGetCount returning "initialization error" means cuInit() did not succeed in that container (missing/mismatched driver or GPU not exposed). Verify nvidia-smi/device passthrough on swx-rdmz-ucx-gpu-02 before attributing further failures to PRs.
• Re-run the job; this failure is unrelated to the PR under test.

Related: PR #11895 (build under test, unrelated to the failure); no existing issue found for the fragile first-line parsing in test_fuzzy_match.py.

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[svc-nvidia-pr-review]

Process/coverage note: the previously reported test-coverage gaps still hold — fragmented transfer with len < 3072 through ucs_memcpy_relaxed() is untested, and the ERMS NT_SOURCE path is never covered with production AMD settings.

[arun-chandran-edarath]

Process/coverage note: the previously reported test-coverage gaps still hold — fragmented transfer with len < 3072 through ucs_memcpy_relaxed() is untested, and the ERMS NT_SOURCE path is never covered with production AMD settings.

The existing guarded nt_buffer_transfer_test() already calls the
dispatcher for every length below 1K and samples the 1K..3K range in
53-byte steps, across all 64 source and destination offsets. For
len < 3072, that dispatcher path intentionally calls libc memcpy().

A second behavioral test through ucs_memcpy_relaxed() cannot distinguish
the dispatcher's libc fallback from the final libc fallback in
ucs_memcpy_relaxed() because both produce the same bytes. It therefore
would not prove which route was taken and would only retest memcpy(). The
configuration regression separately verifies that a finite explicit NT
threshold closes the outer ERMS window by setting
builtin_memcpy_max == 0.

Under production AMD settings, builtin_memcpy_min=inf, so the existing
forced-threshold test is how the generic ERMS source branch is covered on
AMD CI. Production-default ERMS coverage requires an Intel runner.

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[svc-nvidia-pr-review]

Setting builtin_memcpy_max to 0 disables ERMS for every copy, including transfers with total_len < nt_buffer_transfer_min that never reach the NT dispatcher. On Intel the default window is 1K..8M, so UCX_NT_BUFFER_TRANSFER_MIN=<size> silently drops ERMS for all copies below that size, and the user's BUILTIN_MEMCPY_MAX value is discarded. Can we gate the outer window on total_len in ucs_memcpy_relaxed() instead of overwriting the config?

[arun-chandran-edarath]

This is intentional and is the same trade-off discussed in the earlier
NT_BUFFER_TRANSFER_MIN thread.
A finite explicit threshold opts into transfer-level NT selection, so the
outer fragment-length ERMS window is closed to prevent fragments of an
eligible logical transfer from returning before the total_len policy is
evaluated. Explicit NT_BUFFER_TRANSFER_MIN=inf now preserves the configured
BUILTIN_MEMCPY_MAX because no transfer can be NT-eligible.

Gating the outer window on total_len would add another global load and
comparison before every copy in the latency-sensitive Intel ERMS range. I
prefer to keep that path unchanged and apply the explicit opt-in policy during
initialization. The consequence is intentional: a finite explicit NT
threshold disables the outer ERMS window globally, including for
below-threshold transfers.

[svc-nvidia-pr-review]

This skips NT stores for any fragment below 3072 even when total_len is far above nt_dest_threshold; previously fragments above 128 B still used streaming stores. For a transport whose segment size is under 3K the NT_DEST hint becomes a no-op — is that intended?

[arun-chandran-edarath]

Yes, this is intentional. The reworked kernels use 3072 bytes as a
per-fragment crossover, independently of total_len and nt_dest_threshold,
to avoid their vector and streaming-store setup for smaller fragments. A
transport that always segments below 3K therefore uses libc memcpy() for
those fragments even when the logical transfer carries NT_DEST. This is a
deliberate change from the old implementation's 128-byte cutoff.

[svc-nvidia-pr-review]

minor: this 24-line comment repeats what is documented at ucs_cpu_nt_bt_thresh_min() and in the NT_BUFFER_TRANSFER_MIN help string. pls condense to the "why" only (docs/CodeStyle.md: usually one line, document once in the most relevant area).

[arun-chandran-edarath]

This repeats the already resolved
NT_SOURCE documentation thread.
I prefer to keep this documentation because the implicit NT_SOURCE
fallthrough and the interaction between the outer and inner ERMS gates are not
obvious from the code. Fragments below 3072 still reach this function and
return from its initial memcpy() check.

[svc-nvidia-pr-review]

minor: the four 256 B helpers differ only in load/store flavor. Can we generate them with a macro (or pass the load/store choice) to cut ~100 lines of duplication? The same applies to the two kernels, which differ only in stream vs store plus the final sfence.

[arun-chandran-edarath]

This is the same trade-off discussed in the earlier
explicit AVX2 helpers thread.
I prefer to keep the helpers and kernels explicit so the aligned versus
unaligned loads, regular versus streaming stores, and destination fence remain
directly visible and auditable. A macro would reduce source lines but would
not reduce generated code.

[svc-nvidia-pr-review]

minor: ucs_unlikely(src_aligned) hints the wrong way here — (dp-sp) & 31 is 0 whenever both buffers are 64 B-aligned, which is the common case in the copy paths. Also the NT_SOURCE twin at L287 has no hint, so the two kernels disagree.

[arun-chandran-edarath]

The hint is intentional. On the shared-memory AM bcopy path, the
receive-descriptor base is cache-line aligned, but UCP writes its protocol
header first and packs the payload at hdr + 1. Common tag-eager header sizes
are 8, 16, and 24 bytes, so an aligned user source and aligned descriptor base
do not produce relatively aligned payload pointers modulo 32. The
destination-prefix copy advances both pointers equally, leaving the source
unaligned when the streaming body begins.

This preserves the pre-existing AVX2 implementation's
ucs_likely(source_is_unaligned) expectation. The unhinted NT_SOURCE branch
does not make the NT_DEST hint incorrect, so I prefer to retain it.

[svc-nvidia-pr-review]

minor: there is only one implementation, so can we name the kernels ucs_x86_nt_dst_buffer_transfer/ucs_x86_nt_src_buffer_transfer directly and drop these aliases?

[arun-chandran-edarath]

This is the same request as the already resolved
ISA bridge thread.
I prefer to keep the aliases because the stacked AVX-512 follow-up uses them
as the compile-time bridge between AVX2 and AVX-512 while cpu.c keeps
identical dispatcher call sites.

[svc-nvidia-pr-review]

minor: ucs/sys/string.h looks unused — ucs_config_sprintf_memunits() comes from config/parser.h and UCS_MEMUNITS_* from config/types.h.

[arun-chandran-edarath]

ucs/sys/string.h is required here because this file uses
UCS_MEMUNITS_INF, which is defined in that header.
ucs_config_sprintf_memunits() comes from config/parser.h, but
config/types.h does not define UCS_MEMUNITS_INF, so I prefer to keep the
direct include.