[EnricoDeg]

Motivation

Improve handling of large tensor edge cases in convolution

JIRA ID : ROCM-27829

Technical Details

Buffer resource is created with tensor element size as int but tensor byte size is stored as uint. In this way, for 16 bit types, a tensor byte size of exactly 2 GB works. So current limitations are:

• tensor element size is stored in int32_t so max value is 2GB - 1
• tensor element bytes is stored in uint32_t so max value is 4GB - 1.

Convolution fwd supports splitN so if strides are within int range, first splitN is applied. If after splitN, the element size is >= 2GB or the element byte size is > 2GB, the instance is not supported.

Convolution bwd data:

• if the implementation supports splitN, the same workflow is used as convolution fwd
• if the implementation doesn't support splitN, the tensor size is checked when the argument is created and stride_overflow is set. Then this is checked in IsSupportedArgument.

Convolution bwd weight doesn't have splitN support because N is in the GemmK dimension so we always check tensor size when we create the Argument and set stride_overflow. Then this is checked in IsSupportedArgument.

Bugs:

• convolution bwd data xdl cshuffle v1 implementation is using a fast path but the conditions were not exhaustive:
  • num_workgroups_per_Conv_N_ must be 1
  • ConvK % AK1 == 0 && ConvK % BK1 == 0 because GemmK dimension is unmerged into K0 and K1 but K1 is not padded. I decided to not support this case instead of adding padding on the vector dimension which might create performance regressions on the fast path. This is the vector dimension of the descriptor but the actual vector dimension can be lower and this is where the problem comes from.
  • Modify flaky instance on gfx1250

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[therock-pr-bot]

✅ 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_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle.hpp, projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp, projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp, projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v3.hpp, projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp (+10 more); 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]

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

[chris-tsiaousis-hpc]

LGTM! Added some comments/questions

[JP-Fernando]

This PR is correctly scoped and documented. The latest commits fix some earlier issues, already addressed by the other reviewer. Approved with a few comments/questions out of curiosity.

[Copilot]

🟡 Changes recommended

Several size checks can miss oversized tensors or use incorrect data types and descriptor-loop bounds.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Improves convolution handling for large tensors and tightens backward-data fast-path eligibility.

Changes:

• Uses unsigned buffer byte sizes and adds 2 GiB tensor/descriptor checks.
• Propagates overflow detection across backward convolution implementations.
• Restricts an unsafe fast path and adjusts a BF16 forward instance.

File summaries

File	Description
device_grouped_conv_fwd_xdl_instance.hpp	Adjusts BF16 vector widths.
amd_buffer_addressing_builtins.hpp	Uses unsigned resource sizes.
transform_conv_fwd_to_gemm.hpp	Adds element-count limits.
tensor_size_check.hpp	Expands large-tensor checks.
device_grouped_conv_bwd_weight_xdl_waveletmodel_cshuffle_v3.hpp	Adds overflow rejection.
device_grouped_conv_bwd_weight_xdl_cshuffle.hpp	Adds overflow rejection.
device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp	Adds overflow detection.
device_grouped_conv_bwd_weight_wmma_cshuffle.hpp	Adds overflow detection.
device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp	Adds overflow detection.
device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp	Checks accumulator workspace size.
device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp	Checks accumulator workspace size.
device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v3.hpp	Checks all tensor sizes.
device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp	Tightens fast path and descriptor checks.
device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle.hpp	Adds overflow detection.
device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp	Adds descriptor-size validation.

Review details

Suppressed comments (2)

projects/composablekernel/include/ck/tensor_operation/gpu/device/tensor_size_check.hpp:19

• This checks the logical element count, not the addressable element space implied by the supplied strides. These convolution APIs accept explicit stride arrays, so a padded tensor (or a long-index stride above INT32_MAX) can have a resource/descriptor span over the limit while stride_overflow remains false and the stride is later narrowed. Compute the maximum offset from both lengths and strides, or pass an already-built descriptor to this check, and update the callers accordingly.

    for(const auto& l : lengths)
        total *= l;
    long_index_t total_bytes = total * sizeof(DataType);

projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle.hpp:912

• The pointer API repeats the wrong GEMM-alias mapping for external tensors. For heterogeneous convolutions this can mark an oversized weight as safe or reject a valid output; use InDataType, WeiDataType, and OutDataType for the input, weight, and output lengths respectively, and make the corresponding correction in the long-index overload.

        const bool stride_ovf = tensor_exceeds_2gb<ADataType>(a_g_n_c_wis_lengths) ||
                                tensor_exceeds_2gb<BDataType>(b_g_k_c_xs_lengths) ||
                                tensor_exceeds_2gb<CDataType>(e_g_n_k_wos_lengths);

• Files reviewed: 15/15 changed files
• Comments generated: 4
• Review effort level: Balanced

---

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

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-10-05 — Included in TheRock via ROCm/TheRock#8658.