[REVIEW] Addition of gemm_strided_batched - #3106
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe change adds cuBLASLt strided-batched matrix layouts and execution, public Strided-batched GEMM
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
cpp/tests/linalg/gemm_basic.cpp (1)
188-212: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend the batched coverage to the remaining argument combinations.
The test covers only the non-transposed, default-coefficient, host-pointer path. The following cases stay untested and each one exercises distinct code in
matmul_strided_batched:
trans_a/trans_bset to true, which changes the layout row/column swap and thelda/ldbrequirements.- Explicit non-null
alphaandbeta, which verifies that existingCpadding and values are combined correctly.DevicePointerMode = true, which exercises thecoef_wrapperdevice allocation path.- A zero batch stride, which is the cuBLASLt broadcast case.
Do you want me to generate these additional test cases?
As per path instructions: "Tests should use synthetic data, validate numerical results and padding, and cover important edge cases."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/linalg/gemm_basic.cpp` around lines 188 - 212, Extend the batched GEMM coverage around gemm_strided_batched to add synthetic cases for transposed A/B operands with corresponding lda/ldb values, non-null alpha and beta while validating C padding and accumulation, DevicePointerMode enabled to exercise device coefficient allocation, and zero batch strides for broadcast behavior. Assert numerical outputs and untouched padding for each case while preserving the existing baseline test.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp`:
- Around line 378-416: Update matmul_strided_batched to validate positive m, n,
k, and batch_count; require lda and ldc to meet the row requirements for the
configured layouts and ldb to meet the corresponding B-matrix requirement;
validate non-null matrix and scalar pointers and permitted pointer locations
before constructing cuBLASLt layouts. Validate strides for invalid values while
explicitly allowing zero strides for broadcasting, and reject any precondition
failure before forwarding inputs to cuBLASLt.
- Around line 417-432: Update the strided-batched matmul flow around
cublasLtMatmul to create or reuse a valid initialized cublasLtMatmulAlgo_t
instead of passing nullptr. Cache heuristic algorithm results and their
descriptors using strides, batch_count, and compute_type in the key, then
configure and pass workspace when supported while preserving the existing
operation and coefficient parameters.
---
Nitpick comments:
In `@cpp/tests/linalg/gemm_basic.cpp`:
- Around line 188-212: Extend the batched GEMM coverage around
gemm_strided_batched to add synthetic cases for transposed A/B operands with
corresponding lda/ldb values, non-null alpha and beta while validating C padding
and accumulation, DevicePointerMode enabled to exercise device coefficient
allocation, and zero batch strides for broadcast behavior. Assert numerical
outputs and untouched padding for each case while preserving the existing
baseline test.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 821297de-e419-4c57-93c2-210c59d3c380
📒 Files selected for processing (3)
cpp/include/raft/linalg/detail/cublaslt_wrappers.hppcpp/include/raft/linalg/gemm.cuhcpp/tests/linalg/gemm_basic.cpp
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Thanks for implementing this! Let's refine the API (see below)
achirkin
added
feature request
Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/include/raft/linalg/gemm.cuh (2)
286-317: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the Doxygen contract for
gemm_batched.The documentation describes undeclared raw-pointer parameters and template types. Document the 3D mdspan arguments, batch-axis convention, supported layouts, scalar-view behavior, and
compute_type_override.As per path instructions, new public APIs require Doxygen documentation.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/gemm.cuh` around lines 286 - 317, Update the Doxygen block for gemm_batched to match its declared API: document the 3D mdspan A, B, and C arguments, explain the batch-axis convention, list supported layouts, describe scalar-view behavior for alpha and beta, and document compute_type_override. Remove references to undeclared raw-pointer parameters and template types while preserving the public API’s actual semantics.Source: Path instructions
328-354: 🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy liftImplement
gemm_batchedfrom its mdspan parameters.The overload passes undeclared identifiers to
detail::matmul_strided_batched, passes optional scalar objects instead of nullable scalar pointers, and ignorescompute_type_override. Derive dimensions, leading dimensions, element strides, and batch count from the 3D mdspans. Validate extents and supported contiguous layouts before dispatch.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/gemm.cuh` around lines 328 - 354, Implement gemm_batched using its x, y, and z mdspan parameters: validate compatible extents and supported contiguous layouts, derive m, n, k, leading dimensions, batch count, and per-batch strides from the mdspan shapes/strides, and pass the correct underlying matrix data pointers. Convert optional alpha and beta values to nullable scalar pointers expected by detail::matmul_strided_batched, and forward compute_type_override instead of an undeclared compute_type; remove all undeclared identifiers from the dispatch.Sources: Coding guidelines, Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/include/raft/linalg/gemm.cuh`:
- Around line 286-317: Update the Doxygen block for gemm_batched to match its
declared API: document the 3D mdspan A, B, and C arguments, explain the
batch-axis convention, list supported layouts, describe scalar-view behavior for
alpha and beta, and document compute_type_override. Remove references to
undeclared raw-pointer parameters and template types while preserving the public
API’s actual semantics.
- Around line 328-354: Implement gemm_batched using its x, y, and z mdspan
parameters: validate compatible extents and supported contiguous layouts, derive
m, n, k, leading dimensions, batch count, and per-batch strides from the mdspan
shapes/strides, and pass the correct underlying matrix data pointers. Convert
optional alpha and beta values to nullable scalar pointers expected by
detail::matmul_strided_batched, and forward compute_type_override instead of an
undeclared compute_type; remove all undeclared identifiers from the dispatch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d25f523a-0974-4c73-a9a1-527f500b558e
📒 Files selected for processing (1)
cpp/include/raft/linalg/gemm.cuh
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
cpp/tests/linalg/gemm_basic.cpp (1)
206-219: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for non-default alpha and beta in the batched tests.
Both batched tests leave
alphaandbetaempty. Thealpha_val = 2.0fandbeta_val = 3.0fpath throughgemm_batchedtherefore stays untested, and thebeta != 0path never verifies that the existingZcontents are accumulated. The device-scalar pointer mode is also untested for the batched overload.Add at least one batched case with
raft::make_host_scalaralpha and beta, and one withraft::make_device_scalar, following the pattern oftest_gemm_pointer_mode_hostandtest_gemm_pointer_mode_device.As per path instructions,
cpp/REVIEW_GUIDELINES.mdrequires that "tests should exercise numerical correctness, layouts, padding/strides, scalar modes, and relevant edge cases using synthetic data".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/linalg/gemm_basic.cpp` around lines 206 - 219, Extend the batched GEMM coverage in the tests containing gemm_batched to exercise non-default alpha=2.0f and beta=3.0f, verifying accumulation with existing C contents and numerical results. Add cases using raft::make_host_scalar and raft::make_device_scalar, following test_gemm_pointer_mode_host and test_gemm_pointer_mode_device while preserving the existing layout, stride, and padding coverage.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/raft/linalg/gemm.cuh`:
- Around line 377-417: Update both `detail::matmul_strided_batched<kDeviceMode>`
calls in the `z_col_major` and row-major branches to explicitly cast extents and
leading dimensions to `uint64_t`, strides to `int64_t`, and the batch count to
`int32_t`. Before casting `z.extent(0)`, validate it is no greater than
`INT32_MAX` and reject or report out-of-range values using the surrounding API’s
established validation behavior.
---
Nitpick comments:
In `@cpp/tests/linalg/gemm_basic.cpp`:
- Around line 206-219: Extend the batched GEMM coverage in the tests containing
gemm_batched to exercise non-default alpha=2.0f and beta=3.0f, verifying
accumulation with existing C contents and numerical results. Add cases using
raft::make_host_scalar and raft::make_device_scalar, following
test_gemm_pointer_mode_host and test_gemm_pointer_mode_device while preserving
the existing layout, stride, and padding coverage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 98e54e7c-b250-4197-853b-f4816a35fa50
📒 Files selected for processing (2)
cpp/include/raft/linalg/gemm.cuhcpp/tests/linalg/gemm_basic.cpp
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Looks very nice, thanks!
Seems like there's an invalid cublas config in the test, otherwise I think this is ready to merge.
|
/merge |
2 participants
This PR adds the wrapper for strided and batched matrix multiplication with cublasLt requested by @achirkin in the reviews for NVIDIA/cuvs#2352 , where this operation is used for batched row-wise dot products. TF32 support is included.
The associated test covers most of the functionality besides eg transposed inputs, but can be expanded if reviewer(s) prefer.