[MLAS] AVX-512 16-wide Erf kernel and NCHWc reorder transpose for MobileClip-S0 model

[swetha097]

[MLAS] AVX-512 optimizations for MobileClip-S0 FP32 CPU inference

• Added a 16-wide AVX-512 Erf kernel for standalone ONNX Erf operations, replacing the previous 8-wide implementation on AVX-512 hardware.
• Implemented a single-pass 16×16 AVX-512 NCHWc reorder transpose, replacing the SSE2-based 4-wide sub-transpose approach for block-16 data reordering.

The performance numbers taken in STRIX 365 with different thread configurations:

NOTE: The performance numbers were tested on July 31st

STRIX 365 configuration:
AMD Ryzen AI 9 365 (Strix Point) w/ Radeon 880M

Target Model: MobileClip-S0 (FP32, CPU)

Unit tests were added:

test_erf.cpp

• MlasComputeErf vs std::erf within polynomial accuracy tolerance — sweeps buffer lengths straddling the 16-lane boundary to cover both the AVX-512 main
  loop and masked-tail path
• Direct comparison of MlasErfKernelAvx512F vs the base MlasErfKernelFma3 with ≤1 ULP agreement — covers NaN propagation, ±inf, denormals, saturation
• Measured divergence on AVX-512 hardware: 0 ULP (bit-exact); 1 ULP is kept as the cross-microarchitecture contract

test_reorder_input.cpp

• MlasReorderInputNchw (NCHW→NCHWc) vs scalar reference via memcmp
• Sweeps channel counts 1–47: exact 16-channel blocks exercise the new MlasReorderInputNchwBlock16Avx512F fast path; partial blocks exercise the scalar
  tail
• Multiple spatial sizes covered

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Adds AVX-512 optimizations for MobileClip-S0 FP32 CPU inference.

Changes:

• Adds a 16-lane AVX-512 Erf kernel and runtime dispatch.
• Adds AVX-512 NCHW↔NCHWc 16×16 transpose paths.
• Adds Erf and input-reorder correctness tests.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
cmake/onnxruntime_mlas.cmake	Builds the new AVX-512 reorder source.
onnxruntime/core/mlas/lib/gelu_avx512f.cpp	Implements the AVX-512 Erf kernel.
onnxruntime/core/mlas/lib/intrinsics/avx512/reorder_avx512f.cpp	Implements 16×16 reorder transposes.
onnxruntime/core/mlas/lib/mlasi.h	Declares the new kernels.
onnxruntime/core/mlas/lib/platform.cpp	Dispatches Erf to AVX-512.
onnxruntime/core/mlas/lib/reorder.cpp	Uses AVX-512 reorder fast paths.
onnxruntime/test/mlas/unittest/test_erf.cpp	Tests Erf accuracy and kernel equivalence.
onnxruntime/test/mlas/unittest/test_reorder_input.cpp	Tests NCHW input reordering.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Hariharan Seshadri (hariharans29)]

Review: PR #31958 — [MLAS] AVX-512 16-wide Erf kernel and NCHWc reorder transpose for MobileClip-S0 model (head a18a3ab)

Author: @swetha097 (AMD, external — same author as PR #31957), co-author @Manogna-Sree. 4 commits, +486 / −0 across 8 files. CI: 0 / 1 checks OK (external gate). @hariharans29 engaged: re-titled with [MLAS] prefix and just requested Copilot balanced review. CLA pending — bot has prompted. 2 participants.

Verdict: approve on correctness grounds, subject to CLA and full CI. Two independent, well-scoped AVX-512 additions targeting MobileClip-S0 FP32 CPU inference on Strix Point:

1. 16-wide MlasErfKernelAvx512F replaces the 8-wide FMA3 baseline on AVX-512 hardware for standalone ONNX Erf ops.
2. MlasReorderTranspose16x16Avx512F does the NCHW↔NCHWc reorder as a full 16×16 transpose in one pass, replacing the SSE2 4×4 sub-transpose loop when BlockSize == 16.

Test coverage is genuinely thorough — the Erf test asserts both ≤1 ULP vs the base MlasErfKernelFma3 and ≤2e-5 abs vs std::erf across lengths straddling the 16-lane boundary, plus NaN / ±inf / denormal / saturating-magnitude coverage. The reorder test does memcmp bit-exact comparison against a scalar reference across channel counts 1–47 × three spatial sizes.

---

What each piece does

Part 1 — 16-wide AVX-512 Erf

• gelu_avx512f.cpp (+39): new MlasErfKernelAvx512FImpl in an anonymous namespace, with a MLASCALL public wrapper MlasErfKernelAvx512F. Structure:

  while (N >= 16) {
    __m512 X = _mm512_loadu_ps(Input);
    __m512 Result = MlasGeluErfAvx512(X, Constants);
    _mm512_storeu_ps(Output, Result);
    Input += 16; Output += 16; N -= 16;
  }
  if (N > 0) {
    __mmask16 TailMask = __mmask16((1u << N) - 1u);
    __m512 X = _mm512_maskz_loadu_ps(TailMask, Input);
    __m512 Result = MlasGeluErfAvx512(X, Constants);
    _mm512_mask_storeu_ps(Output, TailMask, Result);
  }

• mlasi.h (+1): extern "C" declaration.
• platform.cpp (+1): inside the existing AVX-512F guard (Cpuid7[1] & 0x10000) && (xcr0 & 0xE0) == 0xE0, assign this->ErfKernelRoutine = MlasErfKernelAvx512F. Pre-PR this slot was inherited (FMA3) on AVX-512 hardware.

Part 2 — 16×16 NCHWc reorder transpose

• reorder_avx512f.cpp (+144, new file): MlasReorderTranspose16x16Avx512F — the standard Intel-recommended 16×16 float transpose using unpacklo/hi_ps → unpacklo/hi_pd (via ps↔pd casts) → shuffle_f32x4 at three granularities (32-bit, 64-bit, 128-bit). Two public entry points wrap it:
  • MlasReorderInputNchwBlock16Avx512F(S, D, InputSize) — one full 16-channel NCHW→NCHWc block.
  • MlasReorderOutputNchwBlock16Avx512F(S, D, OutputSize) — one full 16-channel NCHWc→NCHW block.
    Both handle the residual (spatial-tail) with a scalar for c in 0..16 loop.
• reorder.cpp (+26): two #if defined(MLAS_TARGET_AMD64) fast-path insertions in MlasReorderInputNchw and MlasReorderOutputNchwThreaded. When BlockSize == 16 && ChannelsThisIteration == 16 (i.e., a full block, not a trailing partial), route through the AVX-512 helper and continue. Partial trailing blocks still take the SSE2 4×4 path.
• mlasi.h (+22): both entry points declared under #if defined(MLAS_TARGET_AMD64).
• onnxruntime_mlas.cmake (+2): reorder_avx512f.cpp added to both setup_mlas_source_for_windows (mlas_platform_srcs_avx512 — compiled with /arch:AVX512) and the Linux mlas_platform_srcs_avx512f (compiled with -mavx512f).

---

Correctness

Erf kernel

The key question is: what does MlasGeluErfAvx512(X, Constants) actually return? The name suggests "GELU-Erf" (the full x · 0.5 · (1 + erf(x/√2)) formula), which would be wrong for a plain Erf op. If that were true, the new test's math check against std::erf would fail catastrophically at values like erf(3.9) ≈ 0.99999 vs GELU(3.9) ≈ 3.9.

Since the test explicitly checks:

MlasComputeErf(Buffer, Buffer, N);
ASSERT_NEAR(Buffer[i], std::erf(...), 2e-5f);

and passes (empirically 0 ULP vs MlasErfKernelFma3, per author's PR description), MlasGeluErfAvx512 must in fact compute erf(X) directly — the naming refers to "the erf primitive used by the GELU kernel", not the full GELU formula. Reasonable interpretation, and the ULP-level test guards against any future drift. ✓

Nit: consider renaming MlasGeluErfAvx512 to something like MlasErfPolynomialAvx512 (with a follow-up PR) so the shared helper's actual job is unambiguous. Not blocking.

Tail mask correctness: For N ∈ [1, 15], (1u << N) - 1u produces the correct low-N-bits mask (e.g., N=15 → 0x7FFF). ✓ _mm512_maskz_loadu_ps zeros non-mask lanes; erf(0) = 0 is the safe polynomial value, no fault, and _mm512_mask_storeu_ps writes back only the mask-live lanes. Also handles the case where Input/Output may not be a multiple of 16 aligned. ✓

Platform dispatch: platform.cpp assignment is inside the AVX-512F guard that also gates GeluErfKernelRoutine, SiluKernelRoutine, GemmFloatKernel, etc. Consistent. Non-AVX-512 machines continue running MlasErfKernelFma3. ✓

16×16 transpose

The three-level interleave-then-shuffle pattern (unpacklo/hi_ps → unpacklo/hi_pd cast trick → shuffle_f32x4) is the canonical Intel recipe for 16×16 float transpose (equivalent to _MM_TRANSPOSE16x16_PS). Structurally sound.

The definitive check is the memcmp bit-exact test:

MlasReorderInputNchw(Input, Output, Channels, InputSize);
ReferenceReorderInput(Channels, InputSize, Input, OutputReference);
ASSERT_EQ(memcmp(Output, OutputReference, ...), 0);

across Channels ∈ [1, 47] × three spatial sizes. This proves the transpose is bit-identical to the trivial scalar reference Output[p*Bs + c] = Input[c*InputSize + p]. If any of the shuffle constants (0x88, 0xDD) or interleave orderings were wrong, this test would catch it in any run. ✓

Padding-lane invariant: The test seeds OutputReference with -0.5f sentinel and then zero-fills the padding channels [Channels, PaddedInputChannels) in the input. This ensures partial trailing blocks (which take the scalar tail path) correctly zero-fill the padding lanes of the NCHWc block. ✓ Good adversarial choice of sentinel.

Runtime dispatch safety for the reorder fast path

The reorder fast path guards on BlockSize == 16 (via MlasNchwcGetBlockSize()), not on an explicit AVX-512 CPUID check. This is safe because MlasNchwcGetBlockSize() returns 16 only on AVX-512 hardware — on AVX2/SSE2 hardware it returns 8/4/1. So the entry into MlasReorderInputNchwBlock16Avx512F (which unconditionally uses _mm512_* intrinsics) can only happen on AVX-512 machines. ✓

Worth a one-line comment on the fast-path branch making this invariant explicit, so a future maintainer doesn't add a BlockSize == 16 code path from a non-AVX-512 source. Non-blocking.

Symbol linkage: The AVX-512 file is compiled with /arch:AVX512 (Windows) / -mavx512f (Linux) via the cmake source-list, so the symbol is built into the MLAS library on AMD64. Loading a program with AVX-512 opcodes in .text doesn't fault; only executing them would — and we've shown execution is gated. ✓

Padding channel handling

The AVX-512 fast path only fires when InputChannelsThisIteration == 16 (full block). Partial trailing blocks fall through to the SSE2 4×4 path — where the existing correctness for zero-padding trailing channels is preserved. The scalar tail inside MlasReorderInputNchwBlock16Avx512F handles spatial-tail positions (not channel padding). No behavioral change for partial-block inputs. ✓

---

Test coverage

test_erf.cpp — very strong

Level 1 (identity vs base kernel):

• TestMatchesBase(N): sweeps N ∈ {1, 3, 7, 15, 16, 17, 31, 32, 33, 48, 63, 64, 255, 1000} — the boundary-adjacent lengths hit both the main loop and the masked tail.
• Compares MlasErfKernelAvx512F vs MlasErfKernelFma3 (the pre-PR base for AVX-512 machines). Custom UlpDiff helper with correct sign-boundary mapping. Asserts ≤ 1 ULP.
• TestSpecialValuesMatchBase(): NaN (propagation asserted separately from ULP), ±inf, ±denormal_min, ±0, ±1e-30, ±10, ±1e30, ±3.9, ±4.1 (bracketing the saturation region).

Level 2 (math vs std):

• TestMathInPlace(N): runs the kernel in-place (same buffer for input and output), which matches how MobileClip's GELU uses it. Asserts ≤ 2e-5f absolute error vs std::erf. Good — this is a real semantic check, not just kernel-to-kernel.

Gating: Both Level-1 tests early-return if GetMlasPlatform().ErfKernelRoutine != MlasErfKernelAvx512F, so non-AVX-512 hardware skips them. Level 2 runs everywhere via MlasComputeErf.

test_reorder_input.cpp — bit-exact

• Sweeps channel counts 1..47 (covers 1-block, 2-block, 3-block with partial trailing), three spatial sizes: 112×112 (large, spatial multiple of 16 → all main-loop), 15×21 = 315 (spatial not multiple of 16 → main + tail), 11×11 (small odd).
• Reference implementation ReferenceReorderInput is the trivial scalar block[hw*Bs + c] = channel < Channels ? Input[channel*InputSize + hw] : 0.0f.
• memcmp bit-exact comparison.
• Padding-channel and padding-lane invariants explicitly tested via sentinel initialization.

Gate: MlasNchwcGetBlockSize() > 1 — skips test on architectures with no NCHWc support. Correct.

Nit: A parallel test_reorder_output.cpp would round out coverage for MlasReorderOutputNchwThreaded's fast-path. Since the internal transpose helper is shared, indirect coverage exists via correctness of the input path — but a direct output test would be a small addition. Follow-up.

---

Nits (all non-blocking)

1. CLA pending. Bot has prompted @swetha097. Since these are AMD-attributed perf numbers on Strix Point, likely company="AMD" reply. Same CLA gate as PR #31957.
2. CI 0 / 1. Same external-contributor gate. @hariharans29 needs /azp run to unlock Azure Pipelines legs — this is critical here because the AVX-512 legs (Windows GPU CUDA, Windows GPU DML, etc.) actually execute the new kernels on real hardware.
3. MlasGeluErfAvx512 naming ambiguity — see correctness section. Consider renaming the shared primitive in a follow-up so the shared math helper's actual job is not conflated with the top-level GELU wrapper.
4. One-line comment on BlockSize == 16 invariant in the fast-path branches of reorder.cpp would help future maintainers avoid a subtle mis-gating.
5. Add a parallel test_reorder_output.cpp to close the direct-coverage gap for MlasReorderOutputNchwBlock16Avx512F. The shared transpose is exercised via input, so this is defense-in-depth, not blocking.
6. Perf: PR description linked a Strix Point perf-image showing MobileClip-S0 numbers with different thread configurations — full delta values not extractable from fetch output. Deferred to @hariharans29's judgment.
7. File additions to cmake are correct (Windows + Linux both). Standard pattern.
8. 4 commits, mostly polish (unit tests strengthening, review-response cleanup). Will squash cleanly.

---

Merge state

• CI: 0 / 1 — needs /azp run from @hariharans29.
• CLA: pending. Author must reply @microsoft-github-policy-service agree company="AMD" (assuming AMD attribution).
• Approvals: none. Copilot balanced review just requested — pending.
• Reviewers: Copilot AI (in-progress). No human assigned yet.
• Labels: none. MLAS / perf would help retention.

---

Bottom line

Two well-scoped AVX-512 optimizations with strong evidentiary test coverage. The Erf test's ≤ 1 ULP cross-kernel identity check plus ≤ 2e-5 math check catches both dispatch bugs and polynomial divergence. The reorder test's memcmp against a scalar reference across 141 (channels × spatial) combinations catches shuffle-constant bugs bit-exact. Runtime dispatch invariants for the reorder fast path are correct-by-construction.

Action items:

1. Author: reply to CLA bot.
2. @hariharans29: /azp run to unlock AVX-512 CI legs, review Copilot's balanced-review output, then stamp on green.
3. Follow-up (post-merge, non-blocking):
   • Rename MlasGeluErfAvx512 to MlasErfPolynomialAvx512 for clarity.
   • Add test_reorder_output.cpp mirror.
   • Add a comment on the BlockSize == 16 invariant to guard against future mis-gating.

Ready to merge once CLA, /azp run, and one MSFT-side approval are in.

[mirounga]

ARM64 compilation reaches unit tests for this new kernel and fails. Please make sure the build for the kernel test is configured as AVX512-only

[swetha097]

ARM64 compilation reaches unit tests for this new kernel and fails.

swetha097 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree company="Multicoreware"

[swetha097]

ARM64 compilation reaches unit tests for this new kernel and fails. Please make sure the build for the kernel test is configured as AVX512-only

Addressed the comment