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

enzyme/Enzyme/Utils.cpp

@@ -2645,26 +2645,26 @@
   using namespace llvm;
   std::map<BasicBlock *, SmallVector<Instruction *, 1>> maybeBlocks;
   BasicBlock *instBlk = inst->getParent();
   for (auto store : stores) {
     BasicBlock *storeBlk = store->getParent();
     if (instBlk == storeBlk) {
       // if store doesn't come before, exit.
 
       if (store != inst) {
         BasicBlock::const_iterator It = storeBlk->begin();
         for (; &*It != store && &*It != inst; ++It)
           /*empty*/;
         // if inst comes first (e.g. before store) in the
         // block, return true
         if (&*It == inst) {
           results.push_back(store);
         }
       }
       maybeBlocks[storeBlk].push_back(store);
     } else {
       maybeBlocks[storeBlk].push_back(store);
     }
   }
 
   if (maybeBlocks.size() == 0)
     return;
@@ -3925,6 +3925,155 @@
   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<Function>(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;

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<llvm::Value *> trans,
                                   llvm::Value *arg_ld, llvm::Value *dim_1,

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 <stdlib.h>
+#include <string.h>
+
+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;
+}