[cublas] Implement dgmm_batch backend - #756
Open
zjin-lcf wants to merge 3 commits into
Open
Conversation
cuBLAS exposes only the non-batched cublas<t>dgmm (no strided/batched variant exists in any CUDA library), so the cuBLAS backend previously threw unimplemented for all dgmm_batch entry points (issue uxlfoundation#599 -> uxlfoundation#562). Implement dgmm_batch for the cuBLAS backend by looping over cublas<t>dgmm: - Buffer strided, USM strided, and USM group APIs - float, double, complex<float>, complex<double> - Row-major handled by flipping the side and swapping m/n, then delegating to the column-major implementation (matches the rocBLAS backend) Also extend the shared dgmm_batch tests with additional incx values (3 and -1) and group_count values (1 and 10) for more rigorous coverage. Verified on an NVIDIA RTX PRO 6000 Blackwell (sm_120), CUDA 13.3 / cuBLAS 13.6: all DgmmBatch{,Stride}{,Usm} tests pass (col/row major, all four types, run-time and compile-time dispatch). Co-authored-by: Cursor <cursoragent@cursor.com>
Switch dgmm_batch to the cublas<t>dgmm_64 entry points and pass int64_t m/n/lda/incx/ldc directly (no int casts), and drop the 32-bit overflow_check so dimensions beyond 2^31 are supported. Co-authored-by: Cursor <cursoragent@cursor.com>
Switch dgmm_batch (strided buffer, strided USM, and grouped USM) to the ILP64 (_64) rocBLAS entry points and drop the overflow_check calls that capped dimensions below 2^31. Dimensions are now passed straight through as int64_t, matching the cuBLAS backend change. Verified on AMD Instinct MI300A (gfx942): all 24 dgmm_batch CT and RT tests pass. Co-authored-by: Cursor <cursoragent@cursor.com>
sknepper
reviewed
Aug 11, 2026
sknepper
left a comment
sknepper
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the cublas implementation, as well as a nice fix for rocBLAS! One concern I had about writing over an input parameter. Please also check the clang-format check results and adjust the formatting as needed.
| int64_t* ldc, int64_t group_count, int64_t* groupsize, \ | ||
| const std::vector<sycl::event>& dependencies) { \ | ||
| for (int64_t i = 0; i < group_count; i++) \ | ||
| left_right[i] = dgmm_flip_side(left_right[i]); \ |
Contributor
There was a problem hiding this comment.
I'm not sure if this is allowed by the spec. Per https://github.com/uxlfoundation/oneMath/blob/develop/docs/spec/domains/blas/dgmm_batch.rst , left_right is an input parameter, and so shouldn't be modified.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
dgmm_batchwas unimplemented (threwunimplemented) in the cuBLAS backend. As noted in #562, cuBLAS provides only the non-batchedcublas<t>dgmm— there is no strided/batcheddgmmin any CUDA library (verified across cuBLAS, cuBLASLt, cuBLASXt, cuSPARSE headers; onlycublasSdgmm/Ddgmm/Cdgmm/Zdgmmand their_64twins exist). This PR implementsdgmm_batchfor the cuBLAS backend by looping overcublas<t>dgmm.dgmm_batchfloat,double,complex<float>,complex<double>sideand swappingm/n(same approach as the rocBLAS backend)incxvia the standard BLAS start-at-end conventionTests
Extends the shared
dgmm_batchtests for stronger coverage:incx = 3andincx = -1casesgroup_count = 1andgroup_count = 10casesTest plan
Verified on an NVIDIA RTX PRO 6000 Blackwell (sm_120), CUDA 13.3 / cuBLAS 13.6, DPC++:
test_main_blas_ct --gtest_filter=*Dgmm*→ all pass (col/row major, all 4 types)test_main_blas_rt --gtest_filter=*Dgmm*→ all passDgmmBatchStride,DgmmBatchStrideUsm,DgmmBatchUsm(group API)Closes #562.