[BrianHarrisonAMD]

Summary

Fixes #12479.

Two layernorm backward workspace tests pass or fail depending on the compute-unit count of the machine running them. Any PR can be blocked at random, and they are equally broken on develop — the green runs simply landed on machines outside the failing band.

• TestLayernormBwdPlan.GetWorkspaceSizeWithoutOptionalTensorsReturnsNonzero
• TestLayernormPlanBuilder.GetMaxWorkspaceSizeBpropWithoutOptionalTensorsReturnsNonzero

Example failure: run 35807379268, job 107050043339 — Which is: 2408464 vs the expected 16u.

Root cause

The backward workspace is the sum of two terms (LayernormBwdPlan.cpp:148-169, helpers at :344-372):

reqd       = 4 * localSize * multiProcessorCount      // localSize = MAX_LOCAL_SIZE = 1024
isParallel = innerSize <= reqd && innerSize * outerSize > reqd
workspace  = 4 * (parallel ? 2*innerSize*parallelSize : 0)       // parallel reduction buffer
           + 4 * (no mean/inv_variance ? 2*outerSize*stride : 0) // fallback mean/rstd buffer

Both tests want the second term, and assert it as an absolute 16u. That only holds when the first term is zero, which for their shape (outerSize=2, innerSize=150528) means a CU count outside the band [37, 73]. Inside that band the parallel branch is taken and adds 2408448 bytes, which is the failure above. Whether a given run lands inside the band is a property of the machine, not of the change under test.

Change

Assert the difference between the same shape with and without mean/inv_variance. Both graphs have identical dims and strides, so they take the same branch and carry the same parallel reduction buffer — it cancels, leaving the fallback mean/rstd term alone:

EXPECT_EQ(planWithoutStats.getWorkspaceSize(handle) - planWithStats.getWorkspaceSize(handle), 16u);

No CU-dependent expected value, no assumption about which branch the machine takes, and the original shape is kept so both branches still get exercised by whichever machines run the suite. Test-only; no product code changes.

Verification

Full hip_kernel_provider_tests unit suite on the development machine: 531 tests, 524 passed, 7 skipped, 0 failed.

The cancellation was checked with the parallel branch live, using a throwaway probe (not committed) with dims {8,3,224,224} — a shape that satisfies isParallel on that machine's CU count:

withStats=4816896 withoutStats=4816960 delta=64

withStats is the parallel reduction buffer alone (nonzero, so the branch is active) and the delta is exactly 4 * 2 * outerSize * stride = 4 * 2 * 8 * 1 = 64 — the fallback term, independent of the parallel term.

The original CU dependence was reproduced by driving the public statics with fabricated hipDeviceProp_t values:

CU= 16 | parallel=0 ws=16
CU= 27 | parallel=0 ws=16
CU= 48 | parallel=1 ws=2408464
CU= 60 | parallel=1 ws=2408464
CU=104 | parallel=0 ws=16
CU=304 | parallel=0 ws=16

ws=2408464 inside the band matches the linked CI failure byte for byte.

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	PR does not contain code files — Unit Test auto-passed
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[EwanC]

Thanks for tracking down this issue and coming up with a fix.

Patch makes sense to me, to rephrase in my own understanding: The plan needs a workspace in either of the cases a) no provided optional stats tensors b) parallel case which is determined based on device CUs. The test assumption was that the parallel case wouldn't be triggered and only the optional tensors would be used in workspace size calculation, but actually some of the devices stressed by CI did have a number of CUs that hit the parallel case and the extra size requirement threw off the hardcoded 16 workspace size value.

TLDR, LGTM

[adickin-amd]

LGTM

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-09-25 — Included in TheRock via ROCm/TheRock#8472.