[blas] Add half and bfloat16 dot support - #769
Conversation
Expose low-precision dot through runtime and compile-time dispatch with FP32 accumulation in cuBLAS and rocBLAS, and version the expanded backend table ABI. Co-authored-by: Cursor <cursoragent@cursor.com>
ArmPL 26.07 adds bf16 BLAS extensions. Use sbdot so accumulation happens in fp32 with a single rounding to bfloat16, matching the cuBLAS/rocBLAS paths and the test reference; bdot accumulates in bfloat16 and misses that contract. Guarded on _ARMPL_BF16_INTERFACE since older ArmPL has no bf16 BLAS and no version macro. Co-authored-by: Cursor <cursoragent@cursor.com>
melonakos
left a comment
There was a problem hiding this comment.
I went into this expecting to find an ABI problem and instead found it handled properly. One critical coordination issue with #770, and one question that needs a maintainer rather than me.
The ABI handling is correct and complete
Adding function pointers into the middle of function_table_t is normally how you silently break every backend built against the old header. Here it's safe, because oneMath guards it:
if (t->version != SPEC_VERSION)
throw math::specification_mismatch();A stale backend .so is rejected outright rather than loaded with every pointer past the insertion point misaligned. So the mechanism sanctions mid-struct insertion — provided the version is bumped in lockstep, which you did: SPEC_VERSION 1 → 2 in function_table_initializer.hpp, and WRAPPER_VERSION 1 → 2 across all 28 backend wrapper files.
I checked that for completeness rather than taking it on trust. Thirty files in the repo mention WRAPPER_VERSION; you bumped 28. The two you didn't touch are src/dft/backends/backend_wrappers.cxx and src/sparse_blas/backends/backend_wrappers.cxx, which are backend-authoring templates containing the literal placeholder #define WRAPPER_VERSION <Wrapper version number>. Correctly left alone. Nothing missed.
Given the guard, inserting hdot/bfdot beside the other dot variants rather than appending at the end is the right call — it keeps the struct readable, and the version bump is what makes it safe. The two are coupled and you got the coupling right.
Critical: this collides with #770, and getting it wrong is silent
#770 also bumps SPEC_VERSION from 1 to 2, and also inserts into function_table.hpp. I checked: it bumps the same 28 wrapper files to the same version 2.
If both merge as they stand, there are two different table layouts both claiming to be version 2. The t->version != SPEC_VERSION check then passes while the struct offsets disagree — which is precisely the corruption the mechanism exists to prevent, except now with no diagnostic at all. A backend built from one PR's tree loaded against the other's would call the wrong function pointer.
So whichever of these lands second must go to SPEC_VERSION 3 (and bump all 28 wrappers again). That's not something the merge conflict will necessarily force — the conflict is in function_table.hpp and the version line, so it'll likely be visible, but it would be easy to resolve the textual conflict while leaving both at 2.
Please coordinate the two explicitly: pick an order, and have the second PR go to 3. I'd suggest landing this one first since dot is the simpler surface, then rebasing #770 onto it. I'm flagging the same thing on #770.
The numerics are right
This was my other concern and it's addressed. A half-precision dot product accumulated in half precision would be unusable for any real n — about 11 bits of mantissa against an error growing with the reduction length. You avoided that:
cublas_native_func(cublasDotEx, err, handle, n, x_, data_type, incx, y_, data_type,
incy, res_, data_type, CUDA_R_32F);Inputs and output are CUDA_R_16F/CUDA_R_16BF but the execution type is CUDA_R_32F, so the accumulation happens in fp32 and only the final result is narrowed. That's the correct choice, and it's worth a brief comment in the code saying so explicitly — it's the single most important decision in this PR and right now a reader has to know what the last cublasDotEx argument means to see it.
Worth confirming the same holds for the CPU backends you added (netlib_level1.cxx, openblas_level1.cxx, mkl_level1.cxx, armpl_level1.cxx, generic_level1_float.cpp) — since neither Netlib nor OpenBLAS has a native half dot, those paths presumably convert or loop, and they should accumulate in float too. If any of them accumulates in the input type, results will differ noticeably between backends for the same call.
For the maintainers, not for me
This adds dot overloads for sycl::half and oneapi::math::bfloat16 to include/oneapi/math/blas.hxx — the public API surface, not just a backend capability. I don't know whether half/bfloat16 dot is part of the oneMath specification or would be a vendor extension here, and that's a governance call rather than a code-review one.
@sknepper — worth a look. If these aren't spec'd, they may belong behind an extensions header rather than the main BLAS surface, and that decision is much cheaper to make now than after it ships.
Smaller notes
The version scheme is coarse, and this PR exposes that. A BLAS-only table change forces WRAPPER_VERSION bumps in the RNG, DFT, LAPACK and sparse BLAS backends too, because there's a single SPEC_VERSION. Consequence for users: every backend shared library must be rebuilt and redistributed together, or it refuses to load. That's the existing design rather than anything you did, but it deserves a release note, and it's an argument for per-domain versioning at some point.
Your CublasPointerModeGuard and the rocblas_pointer_mode_guard from #764 differ. This one restores unconditionally to CUBLAS_POINTER_MODE_HOST; the rocBLAS one queries and restores the previous mode. Both are correct given the respective defaults, but two guards for the same concept that behave differently is the kind of thing that confuses the next reader. Also, the rocBLAS guard lives in a helper header while this one is local to cublas_level1.cpp — the header is the better home. Same nit as #764: both ignore the status returned by cublasSetPointerMode.
Good handling of the USM result pointer. Deriving result_on_device from sycl::get_pointer_type rather than assuming device memory is the right thing, and it means host and shared allocations work correctly.
Summary
dotoverloads forsycl::halfandbfloat16across runtime and compile-time dispatchcublasDotExandrocblas_dot_ex, including host/device USM results and CUDA 10 fallback handlingTest plan
*Dot*CT/RT suites (one pre-existing H100 CTdotc<double>random-data failure passed on 3 immediate reruns)blas-dot-sycloneMath path on H100 and MI210: FP64, FP32, FP16, and BF16 all PASSgit diff --checkToolchain: DPC++ clang 23.0.0git (
fbf7d1fd2cbf), CUDA toolkit 12.5 / CUDA runtime 13.2, ROCm/HIP 7.1.Made with Cursor