[rocm-libraries] ROCm/rocm-libraries#10547 (commit 235faad)

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.