Skip to content

BLAS: add rules for cblas_?dotc_sub and cblas_?dotu_sub - #3181

Draft
vchuravy wants to merge 1 commit into
mainfrom
vc/blas-complex-dot-sub
Draft

vchuravy wants to merge 1 commit into
mainfrom
vc/blas-complex-dot-sub

Conversation

@vchuravy

@vchuravy vchuravy commented Sep 3, 2026

Copy link
Copy Markdown
Member

Adds derivative rules for the complex dot products in their cblas _sub form (e.g. cblas_zdotc_sub64_ as used by Julia's BLAS.dotc), whose result is written through a pointer passed as the last argument instead of being returned.

primal reverse forward
dotc_sub res = Σ conj(x_i)·y_i dx += conj(dres)·y, dy += dres·x dres = dotc(dx,y) + dotc(x,dy)
dotu_sub res = Σ x_i·y_i dx += dres·y, dy += dres·x dres = dotu(dx,y) + dotu(x,dy)

All c/z variants and suffixes ("", _, 64_, _64_) are recognized; the adjoints are emitted as cblas_?axpy calls, and overwritten inputs are cached via cblas_?copy.

tblgen changes

  • New fpret argument type for a scalar result returned through a pointer. In reverse mode its shadow provides DiffeRet and is reset afterwards (the primal overwrote the result); in forward mode it receives the dual. BlasCalls to *_sub functions get a result alloca appended and their result loaded back, mirroring the existing cublas v2 path.
  • New BConj op (complex conjugate of an fp scalar, identity for real types).
  • Per-pattern supportsComplex bit: the blanket "complex inputs not yet supported in reverse mode" error now only applies to patterns that have not opted in.
  • Complex scalars are passed by pointer under the cblas ABI, so byRefFloat accounts for that; the ABI flags are now emitted by one shared helper for all four generators.
  • Small runtime helpers in Utils.cpp for conjugation and for element-aligned load/store of a blas scalar through a pointer that may be passed as an integer (Julia).

Notes

  • The result pointer is attributed nocapture but deliberately not writeonly: with writeonly, the unused-value analysis considers a local result buffer (e.g. a Julia Ref/alloca) unneeded and replaces it with a fictitious phi, while the primal call writing to it is kept. That looks like a general gap in calculateUnusedValuesInFunction for calls whose only writes go to unneeded memory; left as is here.
  • nocapture is no longer applied to integer-typed vector arguments (invalid IR, hit with Julia-style i64 pointers).

Testing

New lit tests for both functions in forward and reverse mode (test/Enzyme/{ForwardMode,ReverseMode}/blas/cblas_zdot{c,u}_sub64_.ll), including the caching path. Verified locally against LLVM 22; the CHECK lines use regexes tolerant of the typed-pointer / nocapture spelling on older LLVM. Also checked by hand that cblas_cdotc_sub (single precision, 32-bit ints) and the Julia-style form with pointers passed as i64 produce the expected code, and that the existing ddot derivatives are unchanged.

🤖 Generated with Claude Code

https://claude.ai/code/session_01XgNCvikUqowFPsvLWUJqVB

Add derivative rules for the complex dot products in their cblas `_sub`
form (e.g. cblas_zdotc_sub64_), whose result is written through a
pointer passed as the last argument instead of being returned.

  dotc: res = sum(conj(x_i) * y_i)   REV: dx += conj(dres)*y, dy += dres*x
  dotu: res = sum(x_i * y_i)         REV: dx += dres*y,       dy += dres*x
  FWD (both): dres = dot(dx, y) + dot(x, dy)

To express this in the BLAS tblgen:
* new `fpret` argument type for a scalar result returned through a
  pointer; its shadow provides DiffeRet in reverse mode (and is reset,
  since the primal overwrote the result) and receives the dual in
  forward mode. BlasCalls to `*_sub` functions get a result alloca
  appended and their result loaded back, like the cublas v2 path.
* new `BConj` op for the complex conjugate of an fp scalar.
* per-pattern `supportsComplex` bit, so the blanket "complex inputs not
  supported in reverse mode" error only applies to patterns that have
  not been checked for complex inputs.
* complex scalars under the cblas ABI are passed by pointer, so
  byRefFloat now accounts for that (shared helper for all generators).

The result pointer is attributed nocapture but not writeonly: with
writeonly, the unused-value analysis drops a local result buffer
(e.g. a julia Ref) while keeping the primal call writing to it.
nocapture is also no longer applied to non-pointer (integer) vector
arguments, which is invalid IR.

Assisted-by: Claude Code (Fable 5.1)
@vchuravy
vchuravy marked this pull request as draft September 15, 2026 14:03

This branch has not been deployed

No deployments
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