diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp index 4903198de059..c52eb2b97f2a 100644 --- a/enzyme/Enzyme/Utils.cpp +++ b/enzyme/Enzyme/Utils.cpp @@ -3925,6 +3925,155 @@ llvm::Value *to_blas_fp_callconv(IRBuilder<> &B, llvm::Value *V, bool byRef, return allocV; } +Value *emitCublasBeginHostMode(IRBuilder<> &B, Function *called, Value *handle, + IRBuilder<> &entryBuilder) { + auto &M = *B.GetInsertBlock()->getParent()->getParent(); + auto &ctx = M.getContext(); + auto *I32 = Type::getInt32Ty(ctx); + + auto *slot = entryBuilder.CreateAlloca(I32, nullptr, "cublas.pointermode"); + B.CreateStore(ConstantInt::get(I32, CublasPointerModeHost), slot); + + Type *getTys[] = {handle->getType(), slot->getType()}; + B.CreateCall(getOrInsertPerCallingConv(M, called, "cublasGetPointerMode_v2", + FunctionType::get(I32, getTys, false)), + {handle, slot}); + Value *saved = B.CreateLoad(I32, slot, "cublas.savedmode"); + + Type *setTys[] = {handle->getType(), I32}; + B.CreateCall(getOrInsertPerCallingConv(M, called, "cublasSetPointerMode_v2", + FunctionType::get(I32, setTys, false)), + {handle, ConstantInt::get(I32, CublasPointerModeHost)}); + return saved; +} + +void emitCublasEndHostMode(IRBuilder<> &B, Function *called, Value *handle, + Value *savedMode) { + auto &M = *B.GetInsertBlock()->getParent()->getParent(); + auto *I32 = Type::getInt32Ty(M.getContext()); + Type *setTys[] = {handle->getType(), I32}; + B.CreateCall(getOrInsertPerCallingConv(M, called, "cublasSetPointerMode_v2", + FunctionType::get(I32, setTys, false)), + {handle, savedMode}); +} + +/// Emit, once per module, a helper that moves one scalar between a pointer the +/// caller owns and a host buffer. The direction depends on the pointer mode the +/// handle was in, so the branch lives inside the helper rather than splitting +/// the block the derivative is being built into. +static Function *getOrInsertCublasScalarHelper(Module &M, Function *called, + Type *fpTy, bool store) { + auto &ctx = M.getContext(); + auto *I32 = Type::getInt32Ty(ctx); + auto *I8Ptr = getInt8PtrTy(ctx); + auto *VoidTy = Type::getVoidTy(ctx); + uint64_t bytes = M.getDataLayout().getTypeAllocSize(fpTy); + + StringRef vecFn = store ? "cublasSetVector" : "cublasGetVector"; + std::string renamed = getRenamedPerCallingConv(called->getName(), vecFn); + std::string name = + "__enzyme_cublas_scalar_" + std::to_string(bytes) + "_" + renamed; + + // (handle, savedMode, cublasPtr, hostPtr) + Type *argTys[] = {I8Ptr, I32, I8Ptr, I8Ptr}; + FunctionType *FT = FunctionType::get(VoidTy, argTys, false); + auto *F = cast(M.getOrInsertFunction(name, FT).getCallee()); + if (!F->empty()) + return F; + + F->setLinkage(Function::LinkageTypes::InternalLinkage); + F->addFnAttr(Attribute::NoUnwind); + F->addFnAttr(Attribute::AlwaysInline); + + auto *entry = BasicBlock::Create(ctx, "entry", F); + auto *devBB = BasicBlock::Create(ctx, "device", F); + auto *hostBB = BasicBlock::Create(ctx, "host", F); + auto *end = BasicBlock::Create(ctx, "end", F); + + Value *handle = F->arg_begin(); + Value *mode = F->arg_begin() + 1; + Value *cublasPtr = F->arg_begin() + 2; + Value *hostPtr = F->arg_begin() + 3; + + IRBuilder<> EntryB(entry); + EntryB.CreateCondBr( + EntryB.CreateICmpEQ(mode, ConstantInt::get(I32, CublasPointerModeDevice)), + devBB, hostBB); + + { + IRBuilder<> DevB(devBB); + // cublasGetVector(n, elemSize, x, incx, y, incy) copies device -> host and + // cublasSetVector the other way; both take the device side as the pointer + // the caller handed us. + Value *src = store ? hostPtr : cublasPtr; + Value *dst = store ? cublasPtr : hostPtr; + Type *tys[] = {I32, I32, I8Ptr, I32, I8Ptr, I32}; + Value *args[] = {ConstantInt::get(I32, 1), + ConstantInt::get(I32, bytes), + src, + ConstantInt::get(I32, 1), + dst, + ConstantInt::get(I32, 1)}; + DevB.CreateCall(getOrInsertPerCallingConv( + M, called, vecFn, FunctionType::get(I32, tys, false)), + args); + DevB.CreateBr(end); + } + + { + IRBuilder<> HostB(hostBB); + Value *src = store ? hostPtr : cublasPtr; + Value *dst = store ? cublasPtr : hostPtr; + HostB.CreateMemCpy(dst, MaybeAlign(), src, MaybeAlign(), + ConstantInt::get(Type::getInt64Ty(ctx), bytes)); + HostB.CreateBr(end); + } + + IRBuilder<> EndB(end); + EndB.CreateRetVoid(); + (void)handle; + return F; +} + +/// Cast to the opaque pointer the scalar helpers take. +static Value *toHelperPtr(IRBuilder<> &B, Value *V) { + auto *I8Ptr = getInt8PtrTy(B.getContext()); + if (V->getType() == I8Ptr) + return V; + if (V->getType()->isIntegerTy()) + return B.CreateIntToPtr(V, I8Ptr); + return B.CreatePointerCast(V, I8Ptr); +} + +Value *emitCublasLoadScalar(IRBuilder<> &B, Function *called, Value *handle, + Value *savedMode, Value *src, Type *fpTy, + IRBuilder<> &entryBuilder, + llvm::Twine const &name) { + auto &M = *B.GetInsertBlock()->getParent()->getParent(); + auto *host = entryBuilder.CreateAlloca(fpTy, nullptr, "cublas.host." + name); + auto *F = getOrInsertCublasScalarHelper(M, called, fpTy, /*store*/ false); + Value *args[] = {toHelperPtr(B, handle), savedMode, toHelperPtr(B, src), + toHelperPtr(B, host)}; + B.CreateCall(F, args); + if (src->getType()->isPointerTy() && src->getType() != host->getType()) + return B.CreatePointerCast(host, src->getType()); + return host; +} + +void emitCublasStoreScalar(IRBuilder<> &B, Function *called, Value *handle, + Value *savedMode, Value *dst, Value *V, + IRBuilder<> &entryBuilder) { + auto &M = *B.GetInsertBlock()->getParent()->getParent(); + auto *host = + entryBuilder.CreateAlloca(V->getType(), nullptr, "cublas.host.out"); + B.CreateStore(V, host); + auto *F = getOrInsertCublasScalarHelper(M, called, V->getType(), + /*store*/ true); + Value *args[] = {toHelperPtr(B, handle), savedMode, toHelperPtr(B, dst), + toHelperPtr(B, host)}; + B.CreateCall(F, args); +} + Value *is_lower(IRBuilder<> &B, Value *uplo, bool byRef, bool cublas) { if (cublas) { Value *isNormal = nullptr; diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index dd9a0f684867..a9dd37a0a8ab 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -2272,6 +2272,43 @@ llvm::Value *to_blas_fp_callconv(llvm::IRBuilder<> &B, llvm::Value *V, llvm::IRBuilder<> &entryBuilder, llvm::Twine const & = ""); +/// Values of cuBLAS's cublasPointerMode_t. +enum CublasPointerMode { + CublasPointerModeHost = 0, + CublasPointerModeDevice = 1, +}; + +/// cuBLAS reads and writes its scalars -- gemm's alpha and beta, dot's result +/// -- from host or from device memory depending on the handle's pointer mode. +/// That is runtime state: a C caller usually leaves the default host mode in +/// place, while CUDA.jl puts every handle it creates into device mode. The +/// derivative sequence Enzyme emits materializes its scalars on the stack, so +/// instead of duplicating every scalar path it forces the handle to host mode +/// for the duration and copies the caller's own scalars in and out around it. +/// Returns the mode that was in effect, to be handed back to +/// `emitCublasEndHostMode` before the primal call runs. +llvm::Value *emitCublasBeginHostMode(llvm::IRBuilder<> &B, + llvm::Function *called, + llvm::Value *handle, + llvm::IRBuilder<> &entryBuilder); + +void emitCublasEndHostMode(llvm::IRBuilder<> &B, llvm::Function *called, + llvm::Value *handle, llvm::Value *savedMode); + +/// Copy a scalar the caller owns into host memory, returning the host pointer. +/// A no-op copy when the handle was already in host mode. +llvm::Value *emitCublasLoadScalar(llvm::IRBuilder<> &B, llvm::Function *called, + llvm::Value *handle, llvm::Value *savedMode, + llvm::Value *src, llvm::Type *fpTy, + llvm::IRBuilder<> &entryBuilder, + llvm::Twine const & = ""); + +/// Write a host-resident scalar back through a pointer the caller owns. +void emitCublasStoreScalar(llvm::IRBuilder<> &B, llvm::Function *called, + llvm::Value *handle, llvm::Value *savedMode, + llvm::Value *dst, llvm::Value *V, + llvm::IRBuilder<> &entryBuilder); + llvm::Value *get_cached_mat_width(llvm::IRBuilder<> &B, llvm::ArrayRef trans, llvm::Value *arg_ld, llvm::Value *dim_1, diff --git a/enzyme/test/Integration/ForwardMode/cublasdot.cpp b/enzyme/test/Integration/ForwardMode/cublasdot.cpp new file mode 100644 index 000000000000..7e74ab55e38d --- /dev/null +++ b/enzyme/test/Integration/ForwardMode/cublasdot.cpp @@ -0,0 +1,140 @@ +// RUN: if [ %llvmver -ge 12 ]; then %clang++ -fno-exceptions -std=c++11 -O0 %s -S -emit-llvm -o - %loadClangEnzyme | %lli -; fi +// RUN: if [ %llvmver -ge 12 ]; then %clang++ -fno-exceptions -std=c++11 -O1 %s -S -emit-llvm -o - %loadClangEnzyme | %lli -; fi +// RUN: if [ %llvmver -ge 12 ]; then %clang++ -fno-exceptions -std=c++11 -O2 %s -S -emit-llvm -o - %loadClangEnzyme | %lli -; fi +// RUN: if [ %llvmver -ge 12 ]; then %clang++ -fno-exceptions -std=c++11 -O3 %s -S -emit-llvm -o - %loadClangEnzyme | %lli -; fi + +// A cuBLAS _v2 entry point returns its scalar through a trailing pointer +// rather than through the call's return value, and whether that pointer is +// host or device memory depends on the handle's pointer mode -- runtime state +// that a C caller usually leaves at the host default but that CUDA.jl sets to +// device on every handle it creates. +// +// The mock below emulates device memory as an arena and rejects a scalar +// pointer that does not match the current mode, so this test fails if Enzyme +// either drops the tangent or puts it in the wrong address space. +// See https://github.com/EnzymeAD/Enzyme.jl/issues/3442. + +#include "../test_utils.h" + +#include +#include + +extern "C" { + +enum { CUBLAS_POINTER_MODE_HOST = 0, CUBLAS_POINTER_MODE_DEVICE = 1 }; + +struct cublasHandle_t { + int mode; +}; + +// Emulated device memory. A pointer is a device pointer iff it lands here. +static char device_arena[4096]; +static size_t device_used = 0; + +static double *device_alloc(size_t n) { + double *p = (double *)(device_arena + device_used); + device_used += n * sizeof(double); + return p; +} + +static bool is_device_ptr(const void *p) { + return (const char *)p >= device_arena && + (const char *)p < device_arena + sizeof(device_arena); +} + +__attribute__((noinline)) int +cublasGetPointerMode_v2(cublasHandle_t *handle, int *mode) { + *mode = handle->mode; + return 0; +} + +__attribute__((noinline)) int cublasSetPointerMode_v2(cublasHandle_t *handle, + int mode) { + handle->mode = mode; + return 0; +} + +__attribute__((noinline)) int cublasSetVector(int n, int elemSize, + const void *x, int incx, void *y, + int incy) { + // host -> device + if (is_device_ptr(x) || !is_device_ptr(y)) + abort(); + memcpy(y, x, (size_t)n * elemSize); + return 0; +} + +__attribute__((noinline)) int cublasGetVector(int n, int elemSize, + const void *x, int incx, void *y, + int incy) { + // device -> host + if (!is_device_ptr(x) || is_device_ptr(y)) + abort(); + memcpy(y, x, (size_t)n * elemSize); + return 0; +} + +__attribute__((noinline)) int cublasDdot_v2(cublasHandle_t *handle, int n, + const double *x, int incx, + const double *y, int incy, + double *result) { + // The scalar has to live where the handle's pointer mode says it does. + if ((handle->mode == CUBLAS_POINTER_MODE_DEVICE) != is_device_ptr(result)) + abort(); + double res = 0; + for (int i = 0; i < n; i++) + res += x[i * incx] * y[i * incy]; + *result = res; + return 0; +} +} + +__attribute__((noinline)) void my_ddot(cublasHandle_t *handle, int n, + const double *x, const double *y, + double *result) { + cublasDdot_v2(handle, n, x, 1, y, 1, result); +} + +extern "C" double __enzyme_fwddiff(void *, ...); +int enzyme_const; + +int main() { + const int N = 4; + double x[N] = {1.0, 2.0, 3.0, 4.0}; + double y[N] = {5.0, 6.0, 7.0, 8.0}; + double dx[N] = {1.0, 0.0, 0.0, 0.0}; + double dy[N] = {0.0, 1.0, 0.0, 0.0}; + + // d(x . y) = dx . y + x . dy == y[0] + x[1] == 5 + 2 == 7 + const double expected = 7.0; + + // Device pointer mode, as CUDA.jl configures its handles: the result and its + // shadow live in device memory. + { + cublasHandle_t handle; + handle.mode = CUBLAS_POINTER_MODE_DEVICE; + double *res = device_alloc(1); + double *dres = device_alloc(1); + *res = 0; + *dres = 0; + __enzyme_fwddiff((void *)my_ddot, enzyme_const, &handle, enzyme_const, N, x, + dx, y, dy, res, dres); + APPROX_EQ(*dres, expected, 1e-10); + // The handle must be left as it was found. + TEST_EQ(handle.mode, CUBLAS_POINTER_MODE_DEVICE); + } + + // Host pointer mode, the default a C caller sees. + { + cublasHandle_t handle; + handle.mode = CUBLAS_POINTER_MODE_HOST; + double res = 0, dres = 0; + __enzyme_fwddiff((void *)my_ddot, enzyme_const, &handle, enzyme_const, N, x, + dx, y, dy, &res, &dres); + APPROX_EQ(dres, expected, 1e-10); + APPROX_EQ(res, 70.0, 1e-10); + TEST_EQ(handle.mode, CUBLAS_POINTER_MODE_HOST); + } + + return 0; +} diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp index 8442d6903249..d6a2b618ddec 100644 --- a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp @@ -2176,6 +2176,14 @@ void emit_fwd_rewrite_rules(const TGPattern &pattern, raw_ostream &os) { << " \n" << " auto callval = call.getCalledOperand(); \n\n"; + // cuBLAS may be reading and writing its scalars in device memory; run the + // derivative with the handle forced to host mode so the stack slots below + // are what it expects, and stage the caller's scalars around it. + os << " Value *cublas_saved_mode = nullptr;\n" + << " if (cublas)\n" + << " cublas_saved_mode = emitCublasBeginHostMode(Builder2, called, " + "arg_handle, allocationBuilder);\n"; + // just make this const one available now to have less variable name repition os << "Value * const_one = to_blas_callconv(Builder2, " "ConstantInt::get(intType, 1), " @@ -2202,6 +2210,17 @@ void emit_fwd_rewrite_rules(const TGPattern &pattern, raw_ostream &os) { << " ? gutils->invertPointerM(orig_" << name << ", Builder2)\n" << " : nullptr;\n"; os << " }\n"; + os << " if (cublas) {\n" + << " arg_" << name + << " = emitCublasLoadScalar(Builder2, called, arg_handle, " + "cublas_saved_mode, arg_" + << name << ", fpType, allocationBuilder, \"" << name << "\");\n" + << " if (d_" << name << " && !isa(d_" << name << "))\n" + << " d_" << name + << " = emitCublasLoadScalar(Builder2, called, arg_handle, " + "cublas_saved_mode, d_" + << name << ", fpType, allocationBuilder, \"d" << name << "\");\n" + << " }\n"; } } @@ -2243,8 +2262,35 @@ void emit_fwd_rewrite_rules(const TGPattern &pattern, raw_ostream &os) { first = false; } os << ");\n"; - os << " if (!gutils->isConstantValue(&call))\n"; - os << " setDiffe(&call, dres, Builder2);\n"; + if (get_blas_ret_ty(pattern.getName()) == "fpType") { + // A cublas _v2 entry point hands its result back through a trailing + // pointer instead of the call's return value, so that pointer's shadow is + // where the tangent belongs -- setDiffe on the call itself would drop it, + // the call being an inactive status code. + os << " if (cublasv2) {\n" + << " auto orig_ret = call.getArgOperand(" << nameVec.size() + << " + offset);\n" + << " if (!gutils->isConstantValue(orig_ret)) {\n" + << " auto d_ret = gutils->invertPointerM(orig_ret, Builder2);\n" + << " applyChainRule(\n" + << " Builder2,\n" + << " [&](Value *d_ret, Value *dres) {\n" + << " emitCublasStoreScalar(Builder2, called, arg_handle,\n" + << " cublas_saved_mode, d_ret, " + "dres,\n" + << " allocationBuilder);\n" + << " },\n" + << " d_ret, dres);\n" + << " }\n" + << " } else if (!gutils->isConstantValue(&call))\n" + << " setDiffe(&call, dres, Builder2);\n"; + } else { + os << " if (!gutils->isConstantValue(&call))\n"; + os << " setDiffe(&call, dres, Builder2);\n"; + } + os << " if (cublas)\n" + << " emitCublasEndHostMode(Builder2, called, arg_handle, " + "cublas_saved_mode);\n"; os << " }\n"; }