fix(ck): prevent int32 overflow in tensor descriptor element space size by ozturkosu · Pull Request #10547 · ROCm/rocm-libraries

ozturkosu · 2026-08-07T22:27:35Z

JIRA ID : ROCM-29071

Fixes ROCM-29071 (synced with SWSPLAT-48860).

Summary

Fixes a signed 32-bit integer overflow (CWE-190 → CWE-787) in Composable Kernel's runtime tensor-descriptor construction, reported via bug bounty as ROCM-29071.

calculate_element_space_size_impl() (the active path under CK_WORKAROUND_SWDEV_275126) and the fallback lambda in make_naive_tensor_descriptor() both computed:

auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i];

For a runtime (dynamic) descriptor, lengths[i] and strides[i] are index_t (int32), so the multiply is performed in 32-bit before being widened into the long_index_t (int64) accumulator. For tensors where a dimension product exceeds INT32_MAX, the product wraps.

Example (K=C=65537, grouped conv bwd weight):

(65537-1) * 65537 = 4,295,032,832 → wraps to 131,072 (int32)
GetElementSpaceSize() returns 131,073 instead of 4,295,098,369
Workspace allocated ≈ 524 KB, but the kernel writes the full ≈ 17 GB region (sized correctly via accumulate_n<long_index_t>) → out-of-bounds GPU write.

Fix

Widen only the runtime operands to long_index_t before the multiply, via a small helper:

template <typename T>
__host__ __device__ constexpr auto widen_runtime_index_to_long(T v)
{
    if constexpr(is_number_v<T> || is_long_number_v<T>)
        return v;                            // compile-time operand: leave untouched
    else
        return static_cast<long_index_t>(v); // runtime operand: widen before multiply
}

auto acc_new = acc_old + widen_runtime_index_to_long(lengths[i] - Number<1>{})
                       * widen_runtime_index_to_long(strides[i]);

Why not a blanket `static_cast<long_index_t>` on both operands?

An unconditional cast (the first revision of this PR) also widened the compile-time (Number<>) operands, turning the element space size of a fully-static descriptor from a compile-time integral_constant into a runtime long_index_t. That flips TensorDescriptor::IsKnownAtCompileTime() to false, and every static-descriptor consumer gated on it (e.g. threadwise_tensor_slice_transfer, contraction/gemm instances) fails to instantiate — the gfx950 / gfx1201 / Windows math-libs build breaks seen in the previous CI run.

widen_runtime_index_to_long() widens the runtime path (fixing the overflow) while leaving compile-time operands as Number<>, so static descriptors keep a compile-time-constant element space size — identical to the pre-fix type behavior.

Affected path

device_grouped_conv_bwd_weight_xdl_cshuffle.hpp → GetWorkSpaceSize() → GetWorkspaceSizeBytes() → make_naive_tensor_descriptor(...).GetElementSpaceSize(). Any conv-bwd-weight (XDL cshuffle) call with K·C > INT32_MAX.

Notes for reviewers

Both sites are patched so the fix holds regardless of CK_WORKAROUND_SWDEV_275126.
Mirrors the intent of the already-correct ck_tile pattern (ck_tile/core/tensor/tensor_descriptor.hpp, detail::calculate_element_space_size_impl), adapted to preserve CK's compile-time LongNumber<> element-space-size for static descriptors (ck_tile instead clamps to a runtime index_t).

Test plan

Host type-check: a fully-static descriptor keeps a compile-time-constant element space size (IsKnownAtCompileTime() stays true); the runtime path computes 65536 * 65537 = 4,295,032,832 with no int32 wrap.
GPU regression in test/grouped_convnd_bwd_weight/ at K·C > INT32_MAX: reported workspace size matches the 64-bit c_space_size_bytes; kernel run returns hipSuccess.
Audit GetWorkspaceSizeBytes() / c_space_size_bytes callers agree end-to-end.

calculate_element_space_size_impl() and the fallback lambda in make_naive_tensor_descriptor() computed (lengths[i] - 1) * strides[i] in index_t (int32) before widening into the long_index_t accumulator. For tensors where a dimension product exceeds INT32_MAX (e.g. K=C=65537 in grouped conv bwd weight), the multiply wrapped and GetElementSpaceSize() severely undercounted the workspace, while the kernel launch path sized the same region correctly via accumulate_n<long_index_t> -- an overflow -> OOB write mismatch (CWE-190 -> CWE-787). Cast both operands to long_index_t before the multiply at both sites, matching the existing correct pattern in ck_tile/core/tensor/tensor_descriptor.hpp. Fixes ROCM-29071.

therock-pr-bot · 2026-08-07T22:32:01Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	⚠️ Warning	Error: Source/code files changed without an accompanying unit test. Expected: add at least one test file named like `test_<name>.py` / `test_<name>.cpp` (or `<name>_test.`). Current:* code file(s) changed: `projects/composablekernel/include/ck/tensor_description/tensor_descriptor_helper.hpp`, `projects/composablekernel/test/util/unit_tensor_descriptor_element_space_size.cpp`; no test file found
🔎 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 · 2026-08-07T22:32:02Z

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

Copilot

Pull request overview

Fixes 32-bit overflow when calculating Composable Kernel tensor-descriptor element-space sizes.

Changes:

Widens runtime length and stride multiplication to 64 bits.
Updates both compiler-workaround and fallback calculation paths.

Suppressed comments (2)

projects/composablekernel/include/ck/tensor_description/tensor_descriptor_helper.hpp:73

The fallback path has the same compile-time regression: converting Number<> values to scalar long_index_t makes an otherwise static descriptor report IsKnownAtCompileTime() == false, which excludes it from StaticTensor. Preserve the integral-constant result while widening the multiplication.

        auto acc_new = acc_old + static_cast<long_index_t>(lengths[i] - Number<1>{}) *
                                     static_cast<long_index_t>(strides[i]);

projects/composablekernel/include/ck/tensor_description/tensor_descriptor_helper.hpp:30

This overflow fix is not covered by a regression test, despite existing GTest coverage for tensor-descriptor helpers in test/util/unit_tensor_descriptor_functors.cpp. Add the planned runtime {65537, 65537}/{65537, 1} assertion and a fully-static type/IsKnownAtCompileTime() assertion so both the overflow and compile-time descriptor contract are protected.

    auto acc_new = acc_old + static_cast<long_index_t>(lengths[i] - Number<1>{}) *
                                 static_cast<long_index_t>(strides[i]);

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

…overflow (ROCM-29071) The initial ROCM-29071 fix cast both operands to long_index_t unconditionally. For fully-static tensor descriptors this turned the element space size from a compile-time integral_constant into a runtime long_index_t, so TensorDescriptor::IsKnownAtCompileTime() became false and many static-descriptor consumers failed to instantiate (gfx950/gfx1201/Windows math-libs build breaks). Introduce widen_runtime_index_to_long(): widen only runtime index operands to long_index_t before the multiply, leaving Number<>/LongNumber<> operands untouched. Static descriptors keep a compile-time-constant element space size; the runtime path is still int64 and overflow-safe (K*C > INT32_MAX).

Production source should not point to internal JIRA tickets; ticket context belongs in the PR/commit metadata instead.

ozturkosu · 2026-08-11T13:26:37Z

GPU verification (MI300X / gfx942, ROCm 6.4.1)

Built the actual make_naive_tensor_descriptor() from this branch (fixed header overlaid on the CK include tree) and checked GetElementSpaceSize() on both host and device for the reported case (2D {K,C} = {65537,65537}, strides {C,1}):

path	with this fix	old header (control)	expected
host `ess(K=C=65537)`	4295098369 PASS	131073 FAIL (int32 wrap)	4295098369
device `ess(K=C=65537)`	4295098369 PASS	131073 FAIL (int32 wrap)	4295098369
host `ess(boundary 46340²)`	2147441940 PASS	2147441940 PASS (unchanged)	2147441940
static `IsKnownAtCompileTime()`	true (compile-time `static_assert`)	true	true

The old header reports a 131073-element space for a tensor that actually needs 4.29e9 — i.e. the undersized-workspace → OOB-write mechanism, reproduced on device. The fix returns the correct 64-bit value on host and on the gfx942 device, leaves the sub-INT32_MAX boundary case bit-identical, and keeps fully-static descriptors compile-time constant (IsKnownAtCompileTime() stays true, so no static-kernel build regression).

Also confirmed at the type level with -Woverflow -Werror: the flag fires on the old (K-1)*C constant pattern and is clean on the widened one.

One caveat on the PR itself: it doesn't add the committed gtest for the K·C > INT32_MAX case — the -Woverflow idea is a good CI follow-up but isn't wired into the build here.

Add a gtest for make_naive_tensor_descriptor's GetElementSpaceSize() covering the case where (length-1)*stride exceeds INT32_MAX. The test checks the dynamic path on both host and device (must compute in 64-bit), that a sub-INT32_MAX product is unchanged, and that a fully static descriptor keeps IsKnownAtCompileTime() true. Guards the widen fix in tensor_descriptor_helper.hpp against regressions (ROCM-29071).

andriy-ca

Looks good, but can be improved by moving the test in a folder that already exists.

Relocate the GetElementSpaceSize int32-overflow regression test next to the other tensor-descriptor unit tests in test/util (per review), renaming it to the unit_ convention and registering it in test/util/CMakeLists.txt. Drops the separate test/tensor_description subdirectory.

andriy-ca

LGTM!

…overflow

fix(ck): prevent int32 overflow in tensor descriptor element space size (#10547) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : ROCM-29071 Fixes ROCM-29071 (synced with SWSPLAT-48860). ## Summary Fixes a signed 32-bit integer overflow (CWE-190 → CWE-787) in Composable Kernel's runtime tensor-descriptor construction, reported via bug bounty as **ROCM-29071**. `calculate_element_space_size_impl()` (the active path under `CK_WORKAROUND_SWDEV_275126`) and the fallback lambda in `make_naive_tensor_descriptor()` both computed: ```cpp auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i]; ``` For a **runtime** (dynamic) descriptor, `lengths[i]` and `strides[i]` are `index_t` (int32), so the multiply is performed in **32-bit before** being widened into the `long_index_t` (int64) accumulator. For tensors where a dimension product exceeds `INT32_MAX`, the product wraps. **Example (K=C=65537, grouped conv bwd weight):** - `(65537-1) * 65537 = 4,295,032,832` → wraps to `131,072` (int32) - `GetElementSpaceSize()` returns `131,073` instead of `4,295,098,369` - Workspace allocated ≈ 524 KB, but the kernel writes the full ≈ 17 GB region (sized correctly via `accumulate_n<long_index_t>`) → out-of-bounds GPU write. ## Fix Widen **only the runtime operands** to `long_index_t` before the multiply, via a small helper: ```cpp template <typename T> __host__ __device__ constexpr auto widen_runtime_index_to_long(T v) { if constexpr(is_number_v<T> || is_long_number_v<T>) return v; // compile-time operand: leave untouched else return static_cast<long_index_t>(v); // runtime operand: widen before multiply } ``` ```cpp auto acc_new = acc_old + widen_runtime_index_to_long(lengths[i] - Number<1>{}) * widen_runtime_index_to_long(strides[i]); ``` ### Why not a blanket `static_cast<long_index_t>` on both operands? An unconditional cast (the first revision of this PR) also widened the **compile-time** (`Number<>`) operands, turning the element space size of a **fully-static** descriptor from a compile-time `integral_constant` into a runtime `long_index_t`. That flips `TensorDescriptor::IsKnownAtCompileTime()` to `false`, and every static-descriptor consumer gated on it (e.g. `threadwise_tensor_slice_transfer`, contraction/gemm instances) fails to instantiate — the gfx950 / gfx1201 / Windows `math-libs` build breaks seen in the previous CI run. `widen_runtime_index_to_long()` widens the runtime path (fixing the overflow) while leaving compile-time operands as `Number<>`, so static descriptors keep a compile-time-constant element space size — identical to the pre-fix type behavior. ## Affected path `device_grouped_conv_bwd_weight_xdl_cshuffle.hpp` → `GetWorkSpaceSize()` → `GetWorkspaceSizeBytes()` → `make_naive_tensor_descriptor(...).GetElementSpaceSize()`. Any conv-bwd-weight (XDL cshuffle) call with `K·C > INT32_MAX`. ## Notes for reviewers - Both sites are patched so the fix holds regardless of `CK_WORKAROUND_SWDEV_275126`. - Mirrors the intent of the already-correct ck_tile pattern (`ck_tile/core/tensor/tensor_descriptor.hpp`, `detail::calculate_element_space_size_impl`), adapted to preserve CK's compile-time `LongNumber<>` element-space-size for static descriptors (ck_tile instead clamps to a runtime `index_t`). ## Test plan - [x] Host type-check: a fully-static descriptor keeps a compile-time-constant element space size (`IsKnownAtCompileTime()` stays true); the runtime path computes `65536 * 65537 = 4,295,032,832` with no int32 wrap. - [x] GPU regression in `test/grouped_convnd_bwd_weight/` at `K·C > INT32_MAX`: reported workspace size matches the 64-bit `c_space_size_bytes`; kernel run returns `hipSuccess`. - [ ] Audit `GetWorkspaceSizeBytes()` / `c_space_size_bytes` callers agree end-to-end.

fix(ck): prevent int32 overflow in tensor descriptor element space size JIRA ID : ROCM-29071 Fixes ROCM-29071 (synced with SWSPLAT-48860). ## Summary Fixes a signed 32-bit integer overflow (CWE-190 → CWE-787) in Composable Kernel's runtime tensor-descriptor construction, reported via bug bounty as **ROCM-29071**. `calculate_element_space_size_impl()` (the active path under `CK_WORKAROUND_SWDEV_275126`) and the fallback lambda in `make_naive_tensor_descriptor()` both computed: ```cpp auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i]; ``` For a **runtime** (dynamic) descriptor, `lengths[i]` and `strides[i]` are `index_t` (int32), so the multiply is performed in **32-bit before** being widened into the `long_index_t` (int64) accumulator. For tensors where a dimension product exceeds `INT32_MAX`, the product wraps. **Example (K=C=65537, grouped conv bwd weight):** - `(65537-1) * 65537 = 4,295,032,832` → wraps to `131,072` (int32) - `GetElementSpaceSize()` returns `131,073` instead of `4,295,098,369` - Workspace allocated ≈ 524 KB, but the kernel writes the full ≈ 17 GB region (sized correctly via `accumulate_n<long_index_t>`) → out-of-bounds GPU write. ## Fix Widen **only the runtime operands** to `long_index_t` before the multiply, via a small helper: ```cpp template <typename T> __host__ __device__ constexpr auto widen_runtime_index_to_long(T v) { if constexpr(is_number_v<T> || is_long_number_v<T>) return v; // compile-time operand: leave untouched else return static_cast<long_index_t>(v); // runtime operand: widen before multiply } ``` ```cpp auto acc_new = acc_old + widen_runtime_index_to_long(lengths[i] - Number<1>{}) * widen_runtime_index_to_long(strides[i]); ``` ### Why not a blanket `static_cast<long_index_t>` on both operands? An unconditional cast (the first revision of this PR) also widened the **compile-time** (`Number<>`) operands, turning the element space size of a **fully-static** descriptor from a compile-time `integral_constant` into a runtime `long_index_t`. That flips `TensorDescriptor::IsKnownAtCompileTime()` to `false`, and every static-descriptor consumer gated on it (e.g. `threadwise_tensor_slice_transfer`, contraction/gemm instances) fails to instantiate — the gfx950 / gfx1201 / Windows `math-libs` build breaks seen in the previous CI run. `widen_runtime_index_to_long()` widens the runtime path (fixing the overflow) while leaving compile-time operands as `Number<>`, so static descriptors keep a compile-time-constant element space size — identical to the pre-fix type behavior. ## Affected path `device_grouped_conv_bwd_weight_xdl_cshuffle.hpp` → `GetWorkSpaceSize()` → `GetWorkspaceSizeBytes()` → `make_naive_tensor_descriptor(...).GetElementSpaceSize()`. Any conv-bwd-weight (XDL cshuffle) call with `K·C > INT32_MAX`. ## Notes for reviewers - Both sites are patched so the fix holds regardless of `CK_WORKAROUND_SWDEV_275126`. - Mirrors the intent of the already-correct ck_tile pattern (`ck_tile/core/tensor/tensor_descriptor.hpp`, `detail::calculate_element_space_size_impl`), adapted to preserve CK's compile-time `LongNumber<>` element-space-size for static descriptors (ck_tile instead clamps to a runtime `index_t`). ## Test plan - [x] Host type-check: a fully-static descriptor keeps a compile-time-constant element space size (`IsKnownAtCompileTime()` stays true); the runtime path computes `65536 * 65537 = 4,295,032,832` with no int32 wrap. - [x] GPU regression in `test/grouped_convnd_bwd_weight/` at `K·C > INT32_MAX`: reported workspace size matches the 64-bit `c_space_size_bytes`; kernel run returns `hipSuccess`. - [ ] Audit `GetWorkspaceSizeBytes()` / `c_space_size_bytes` callers agree end-to-end. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: muozturk <Osman.Ozturk@amd.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>

ozturkosu mentioned this pull request Aug 7, 2026

fix(ck): prevent int32 overflow in tensor descriptor element space size (ROCM-29071) #10545

Closed

3 tasks

github-actions Bot added the project: composablekernel label Aug 7, 2026

ozturkosu self-assigned this Aug 7, 2026

therock-pr-bot Bot added the Not ready to Review label Aug 7, 2026

assistant-librarian Bot added the organization: ROCm label Aug 7, 2026

ozturkosu requested a balanced review from Copilot August 8, 2026 02:27

Copilot started reviewing on behalf of ozturkosu August 8, 2026 02:28 View session

Copilot AI reviewed Aug 8, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck/tensor_description/tensor_descriptor_helper.hpp Outdated

therock-pr-bot Bot removed the Not ready to Review label Aug 8, 2026

ozturkosu changed the title ~~fix(ck): prevent int32 overflow in tensor descriptor element space size (ROCM-29071)~~ fix(ck): prevent int32 overflow in tensor descriptor element space size Aug 8, 2026

ozturkosu marked this pull request as ready for review August 8, 2026 03:40

ozturkosu requested a review from a team as a code owner August 8, 2026 03:40

ozturkosu marked this pull request as draft August 10, 2026 08:09

ozturkosu marked this pull request as ready for review August 11, 2026 11:45

docs(ck): remove JIRA ticket reference from source comments

055a8b5

Production source should not point to internal JIRA tickets; ticket context belongs in the PR/commit metadata instead.

andriy-ca requested changes Aug 11, 2026

View reviewed changes

Comment thread projects/composablekernel/test/util/unit_tensor_descriptor_element_space_size.cpp

andriy-ca approved these changes Aug 12, 2026

View reviewed changes

ozturkosu enabled auto-merge (squash) August 12, 2026 14:32

ozturkosu and others added 4 commits August 12, 2026 08:33

Merge branch 'develop' into users/muozturk/ck/tensor-description-int-…

9f9cf03

…overflow

Merge branch 'develop' into users/muozturk/ck/tensor-description-int-…

0441ff1

…overflow

Merge branch 'develop' into users/muozturk/ck/tensor-description-int-…

af142bf

…overflow

Merge branch 'develop' into users/muozturk/ck/tensor-description-int-…

3e4c32f

…overflow

ammallya disabled auto-merge August 14, 2026 21:58

ammallya merged commit 235faad into develop Aug 14, 2026
21 checks passed

ammallya deleted the users/muozturk/ck/tensor-description-int-overflow branch August 14, 2026 21:58

ozturkosu mentioned this pull request Aug 18, 2026

Revert "fix(ck): prevent int32 overflow in tensor descriptor element space size" (#10547) #10947

Closed

Navigation Menu

Uh oh!

fix(ck): prevent int32 overflow in tensor descriptor element space size - #10547

fix(ck): prevent int32 overflow in tensor descriptor element space size#10547
ammallya merged 9 commits into
developfrom
users/muozturk/ck/tensor-description-int-overflow

ozturkosu commented Aug 7, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Aug 7, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Aug 7, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

ozturkosu commented Aug 11, 2026

Uh oh!

andriy-ca left a comment

Uh oh!

Uh oh!

andriy-ca left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

ozturkosu commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Fix

Why not a blanket static_cast<long_index_t> on both operands?

Affected path

Notes for reviewers

Test plan

Uh oh!

therock-pr-bot Bot commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

ozturkosu commented Aug 11, 2026

GPU verification (MI300X / gfx942, ROCm 6.4.1)

Uh oh!

andriy-ca left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

andriy-ca left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ozturkosu commented Aug 7, 2026 •

edited

Loading

Why not a blanket `static_cast<long_index_t>` on both operands?

therock-pr-bot Bot commented Aug 7, 2026 •

edited

Loading

therock-pr-bot Bot commented Aug 7, 2026 •

edited

Loading