Deliver the forward-mode tangent of a cuBLAS _v2 scalar result - #3113
Open
wsmoses wants to merge 1 commit into
Open
Deliver the forward-mode tangent of a cuBLAS _v2 scalar result#3113wsmoses wants to merge 1 commit into
wsmoses wants to merge 1 commit into
Conversation
A cuBLAS _v2 entry point hands its result back through a trailing pointer rather than through the call's return value. emit_fwd_rewrite_rules finished with setDiffe(&call, dres), and for these entry points the call itself is an inactive cublasStatus_t, so the tangent was computed and then dropped: forward-mode cublasDdot_v2 returned zero. It now writes dres through the shadow of the result argument. Doing that correctly means confronting the pointer mode. cuBLAS reads and writes its scalars -- gemm's alpha and beta, dot's result -- from host or from device memory depending on cublasSetPointerMode_v2, which is runtime state. A C caller usually leaves the host default in place, but CUDA.jl puts every handle it creates into device mode, where the stack slots the derivative code materializes its scalars in are not addresses cuBLAS can use. Rather than duplicate every scalar path, the derivative sequence now runs with the handle temporarily forced to host mode and the caller's own scalars copied in and out around it -- the same bracketing cuBLAS.jl itself does. This needs no device allocation, and only uses entry points that libcublas exports (cublasGet/SetPointerMode_v2 and cublasSet/GetVector); cudaMemcpy and cuMemcpy are not reachable from a libcublas handle. The branch on the original mode lives inside a module-local helper so that generated code never has to split the block it is being built into. Reverse mode is untouched. The new test emulates device memory as an arena and rejects a scalar pointer that does not match the handle's mode, so it fails if the tangent is dropped or lands in the wrong address space. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6a8BiaAKpUTgP86mQ9ZKP
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
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.
A cuBLAS
_v2entry point hands its result back through a trailing pointerrather than through the call's return value.
emit_fwd_rewrite_rulesfinishedwith
setDiffe(&call, dres), and for these entry points the call itself is aninactive
cublasStatus_t, so the tangent was computed and then discarded —forward-mode
cublasDdot_v2returned zero. Before:It now writes
dresthrough the shadow of the result argument.The pointer mode
Doing that correctly means confronting the pointer mode. cuBLAS reads and writes
its scalars — gemm's
alphaandbeta, dot's result — from host or from devicememory depending on
cublasSetPointerMode_v2, which is runtime state. A Ccaller usually leaves the host default in place, but CUDA.jl puts every handle
it creates into device mode, where the stack slots the derivative materializes
its scalars in are not addresses cuBLAS can use.
Rather than duplicate every scalar path, the derivative sequence now runs with
the handle temporarily forced to host mode and the caller's own scalars copied
in and out around it — the same bracketing cuBLAS.jl itself does. This needs no
device allocation and uses only entry points libcublas exports
(
cublasGet/SetPointerMode_v2,cublasSet/GetVector);cudaMemcpyandcuMemcpyare not reachable from a libcublas handle. The branch on theoriginal mode lives in a module-local helper, so generated code never has to
split the block it is being built into.
Reverse mode is untouched and still assumes host pointer mode.
Testing
Integration/ForwardMode/cublasdot.cppemulates device memory as an arena andrejects a scalar pointer that does not match the handle's mode, so it fails if
the tangent is dropped or lands in the wrong address space. It covers both
pointer modes and checks the handle is left as it was found.
The pre-existing
ReverseMode/cublas.cppand the forward/reverseblas.cppintegration tests all still pass. Full unit suite on LLVM 16: 1189 pass, 11
xfail, 1 fail — that failure is
ReverseModeVector/partial_int_window.ll, whichfails identically on unmodified
mainwith this local LLVM 16 build. Format(clang-format 16) and
check_emission_order.pyclean.Split out of #3102 at review request. Independent of the CUDA transfer
derivatives there — the two touch disjoint code — so this can land on its own.
🤖 Generated with Claude Code
https://claude.ai/code/session_01R6a8BiaAKpUTgP86mQ9ZKP