Skip to content

Deliver the forward-mode tangent of a cuBLAS _v2 scalar result - #3113

Open
wsmoses wants to merge 1 commit into
mainfrom
cublas-fwd-scalar
Open

Deliver the forward-mode tangent of a cuBLAS _v2 scalar result#3113
wsmoses wants to merge 1 commit into
mainfrom
cublas-fwd-scalar

Conversation

@wsmoses

@wsmoses wsmoses commented Aug 9, 2026

Copy link
Copy Markdown
Member

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 discarded —
forward-mode cublasDdot_v2 returned zero. Before:

define internal fastcc void @fwddiffe_my_ddot(...) {
  %1 = call i32 @cublasDdot_v2(..., ptr %"x'", ..., ptr %y, ..., ptr %0)
  %3 = call i32 @cublasDdot_v2(..., ptr %x, ..., ptr %"y'", ..., ptr %2)
  ret void                                        ; tangent discarded
}

It now writes dres through 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 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 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); cudaMemcpy and
cuMemcpy are not reachable from a libcublas handle. The branch on the
original 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.cpp 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. It covers both
pointer modes and checks the handle is left as it was found.

The pre-existing ReverseMode/cublas.cpp and the forward/reverse blas.cpp
integration tests all still pass. Full unit suite on LLVM 16: 1189 pass, 11
xfail, 1 fail — that failure is ReverseModeVector/partial_int_window.ll, which
fails identically on unmodified main with this local LLVM 16 build. Format
(clang-format 16) and check_emission_order.py clean.

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

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant