[ROCm] redirect C64/C128 matmuls to rocBLAS at runtime#924
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| auto run = [&](auto alpha, auto beta) -> absl::Status { | ||
| if (batch_size != 1) { | ||
| return blas->DoBlasGemmStridedBatched( | ||
| stream, trans_a, trans_b, m, n, k, dtype, &alpha, a_buf, lda, | ||
| lhs.batch_stride, b_buf, ldb, rhs.batch_stride, &beta, &out_buf, ldc, | ||
| out.batch_stride, static_cast<int>(batch_size), engine_options, | ||
| context); | ||
| } |
There was a problem hiding this comment.
question: When batch_size is 0 (both lhs and rhs have batch_size == 0), the condition batch_size != 1 is true, so this takes the strided-batched path with batch_count == 0. Most BLAS implementations handle this as a no-op, but is that the intended behavior? Would it be clearer to guard with batch_size > 1 (or batch_size >= 2) and treat 0 like 1?
Review SummaryClean, well-scoped approach: the GemmRewriter adds complex type entries so hipBLASLt custom calls are emitted as usual, and 6 inline comments posted — mostly minor (TODO typo, narrowing casts, unused One broader note: the PR has good local test coverage (222 JAX tests, 7 bazel tests), but no new unit tests are included in the diff. An HLO-level test exercising the C64/C128 rocBLAS fallback path would help prevent regressions. |
github-actions
Bot
removed
the
claude-review
magaonka-amd
force-pushed
the
feature/complex-gemm-rocblas-redirect
branch
from
9008301 to
2f9b078
Compare
After XLA removed the cublas/rocblas backend, ROCm has no way to compute C64/C128 GEMMs because hipBLASLt does not support complex GEMMs yet. This PR redirects complex GEMM calls to rocBLAS at runtime. hipBLASLt support for complex GEMMs is in progress in ROCm; once that lands, I'll version-gate this code. Testing: - Tested 222 complex GEMM tests in JAX pytest and 7 bazel tests using complex GEMM. - From the XLA side only tested locally. Once this PR merges we will need to re-enable the skipped complex GEMM tests in XLA.
magaonka-amd
force-pushed
the
feature/complex-gemm-rocblas-redirect
branch
from
2f9b078 to
61abdf0
Compare
| absl::StatusOr<std::vector<MatmulAlgorithm>> GetAlgorithms( | ||
| size_t /*max_algorithm_count*/, | ||
| size_t /*max_workspace_size*/) const override { | ||
| return std::vector<MatmulAlgorithm>{MatmulAlgorithm{std::any{}, 0}}; |
There was a problem hiding this comment.
This is nice, but there is an open question if it is better to actulay expose rocblas algorithms and workspace.
| }; | ||
| blas::Transpose trans_a = transpose_of(lhs), trans_b = transpose_of(rhs); | ||
|
|
||
| const uint64_t m = out.num_rows, n = out.num_cols, k = lhs.num_cols; |
There was a problem hiding this comment.
oh we need to be careful here. Do we have tests with different values for trans_a/b flags?
I remember for hipblaslt C++ API, I had to do the following:
MatmulPlan::Config internal_cfg {
.m = output_layout.num_rows,
.n = output_layout.num_cols,
.k = (hip_trans_a == HIPBLAS_OP_N ? lhs_layout.num_cols
: lhs_layout.num_rows),
.batch_count = lhs_layout.batch_size,
.lda = lhs_layout.leading_dim_stride,
.ldb = rhs_layout.leading_dim_stride,
.ldc = c_layout.leading_dim_stride,
.ldd = output_layout.leading_dim_stride,
.strideA = lhs_layout.batch_stride,
.strideB = rhs_layout.batch_stride,
.strideC = c_layout.batch_stride,
.strideD = output_layout.batch_stride,
.epilogue = gemm_epi
};
There was a problem hiding this comment.
it looks like this change you mentioned operates on raw layout ?
In our case I do : bool must_swap_operands = gpu::MakeOutputColumnMajor(lhs, rhs, out);
https://github.com/ROCm/xla/pull/924/changes#diff-4c837b4660f5e0bc5378a3f531172180e0aaf62a38159defe17d6f3be6a17ef9R580
in that case k = lhs.num_cols; still okay I think.
is my understanding correct?
There was a problem hiding this comment.
@magaonka-amd I didn't realise Pavel had comments on this, seems it's not resolved yet and your PR is already merged to upstream #924
@pemeliya could you check whether any new UT, please?
There was a problem hiding this comment.
Oh I did not notice that we compute trans_a/b based on the memory order:
auto transpose_of = [](const gpu::MatrixLayout& l) {
return l.order == gpu::MatrixLayout::Order::kColumnMajor
? blas::Transpose::kNoTranspose
: blas::Transpose::kTranspose;
};
blas::Transpose trans_a = transpose_of(lhs), trans_b = transpose_of(rhs);
For a conventional hipblaslt execution path, trans_a/trans_b are provided by the layouts:
auto trans_a = lhs_layout.transpose, trans_b = rhs_layout.transpose;
I cannot say whether these two are equivalent or not here..
Maybe it makes sense to create some UT with C64 dots and non-square matrices, so that the transpose matters and try to get all 4 trans_a/trans_b combinations compiled to see if it works correctly
|
change approved upstream closing this |
i-chaochen
left a comment
i-chaochen
left a comment
There was a problem hiding this comment.
also realise one thing, I guess we need to enable these disabled c64/c128 UTs https://github.com/openxla/xla/pull/43029/changes ?
cc @alekstheod
|
@i-chaochen yes we need to reenable these tests in both xla and jax bazel side. @pemeliya I'll explore adding new tests as you suggested , it may take couple more days to get on that . right now working some stuff with JAX CI ,. will finish that and pick this up. thanks for the suggestions. |
4 participants
After XLA removed the cublas/rocblas backend, XLA has no way to compute C64/C128 GEMMs because hipBLASLt does not support complex GEMMs yet.
This PR redirects complex GEMM calls to rocBLAS at runtime. hipBLASLt support for complex GEMMs is in progress in ROCm; once that lands, I'll version-gate this code.
Testing