BLAS: add rules for cblas_?dotc_sub and cblas_?dotu_sub

enzyme/Enzyme/BlasDerivatives.td

@@ -12,6 +12,9 @@ def len   : BLASType<1, 0>; // num of elements
 def hbw   : BLASType<1, 0>; // half matrix-bandwith}
 def bw    : BLASType<1, 0>; // matrix-bandwith}
 def fp    : BLASType<1, 1>;  // floating point
+// pointer through which a scalar fp result is returned, taking the place of
+// the return value (e.g. the last arg of cblas_zdotc_sub)
+def fpret : BLASType<1, 1>;
 
 // packed array of size ( n * (n  + 1) / 2 )
 class ap<list<string> _args> : BLASType<1, 1> {
@@ -48,6 +51,8 @@ class CallBlasPattern<dag patternToMatch, list<string> mutables, list<BLASType>
   list<dag> ArgDerivatives = resultOps;
   list<string> mutable = mutables;
   dag ArgDuals = forwardOps;
+  // Whether the reverse-mode rules are valid for complex (c, z) inputs.
+  bit supportsComplex = 0;
 }
 
 
@@ -56,6 +61,7 @@ def Rows : MagicInst; // given a transpose, normal rows, normal cols get the tru
 def Concat : MagicInst;
 
 def ShadowNoInc : MagicInst;
+def BConj : MagicInst; // complex conjugate of an fp scalar (identity for reals)
 
 class Binop<string _s, list<string> _tys> {
   string s = _s;
@@ -245,6 +251,38 @@ def dot : CallBlasPattern<(Op $n, $x, $incx, $y, $incy),
                   (FAdd<""> (BlasCall<"dot"> $n, (Shadow $x), $y), (BlasCall<"dot"> $n, $x, (Shadow $y)))
                   >;
 
+// res = x^H * y = sum(conj(x_i) * y_i), written through the res pointer
+// (the cblas `_sub` calling convention used for the complex dot products).
+// REV: dx += conj(dres) * y ; dy += dres * x
+// FWD: dres = dotc(dx, y) + dotc(x, dy)
+def dotc_sub : CallBlasPattern<(Op $n, $x, $incx, $y, $incy, $res),
+                  ["res"],[len, vinc<["n"]>, vinc<["n"]>, fpret],
+                  [
+                  (BlasCall<"axpy"> $n, (BConj DiffeRet), $y, (Shadow $x)),
+                  (BlasCall<"axpy"> $n, DiffeRet, $x, (Shadow $y)),
+                  (InactiveArg) // res is overwritten; its shadow is consumed as DiffeRet
+                  ],
+                  (FAdd<""> (BlasCall<"dotc_sub"> $n, (Shadow $x), $y), (BlasCall<"dotc_sub"> $n, $x, (Shadow $y)))
+                  > {
+  let supportsComplex = 1;
+}
+
+// res = x^T * y = sum(x_i * y_i), written through the res pointer
+// (unconjugated complex dot product, cblas `_sub` calling convention).
+// REV: dx += dres * y ; dy += dres * x
+// FWD: dres = dotu(dx, y) + dotu(x, dy)
+def dotu_sub : CallBlasPattern<(Op $n, $x, $incx, $y, $incy, $res),
+                  ["res"],[len, vinc<["n"]>, vinc<["n"]>, fpret],
+                  [
+                  (BlasCall<"axpy"> $n, DiffeRet, $y, (Shadow $x)),
+                  (BlasCall<"axpy"> $n, DiffeRet, $x, (Shadow $y)),
+                  (InactiveArg) // res is overwritten; its shadow is consumed as DiffeRet
+                  ],
+                  (FAdd<""> (BlasCall<"dotu_sub"> $n, (Shadow $x), $y), (BlasCall<"dotu_sub"> $n, $x, (Shadow $y)))
+                  > {
+  let supportsComplex = 1;
+}
+
 def nrm2 : CallBlasPattern<(Op $n, $x, $incx),
                   [],[len, vinc<["n"]>],
                   [

enzyme/Enzyme/Utils.cpp

@@ -2648,26 +2648,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;
@@ -3695,9 +3695,10 @@
 #endif
 {
   const char *extractable[] = {
-      "dot",   "scal",  "axpy",  "gemv",  "gemm",  "spmv", "syrk",  "nrm2",
-      "trmm",  "trmv",  "symm",  "potrf", "potrs", "copy", "spmv",  "syr2k",
-      "potrs", "getrf", "getrs", "trtrs", "getri", "symv", "lacpy", "trsv",
+      "dot",  "scal",  "axpy",  "gemv",     "gemm",  "spmv",  "syrk",
+      "nrm2", "trmm",  "trmv",  "symm",     "potrf", "potrs", "copy",
+      "spmv", "syr2k", "potrs", "getrf",    "getrs", "trtrs", "getri",
+      "symv", "lacpy", "trsv",  "dotc_sub", "dotu_sub",
   };
   const char *floatType[] = {"s", "d", "c", "z"};
   const char *prefixes[] = {"" /*Fortran*/, "cblas_"};
@@ -4217,6 +4218,42 @@
   return B.CreateLoad(intType, V);
 }
 
+llvm::Value *complex_conjugate(llvm::IRBuilder<> &B, llvm::Value *V) {
+  auto VT = dyn_cast<VectorType>(V->getType());
+  if (!VT)
+    return V;
+  assert(VT->getElementCount().getKnownMinValue() == 2 &&
+         "complex blas scalars are 2-element vectors");
+  Value *im = B.CreateExtractElement(V, (uint64_t)1);
+  return B.CreateInsertElement(V, B.CreateFNeg(im), (uint64_t)1);
+}
+
+static llvm::Value *blas_scalar_ptr(llvm::IRBuilder<> &B, llvm::Type *fpType,
+                                    llvm::Value *ptr) {
+  if (ptr->getType()->isIntegerTy())
+    return B.CreateIntToPtr(ptr, getUnqual(fpType));
+  return B.CreatePointerCast(
+      ptr, getPointerType(
+               fpType, cast<PointerType>(ptr->getType())->getAddressSpace()));
+}
+
+static llvm::Align blas_scalar_align(llvm::IRBuilder<> &B, llvm::Type *fpType) {
+  auto &DL = B.GetInsertBlock()->getModule()->getDataLayout();
+  return DL.getABITypeAlign(fpType->getScalarType());
+}
+
+llvm::Value *load_blas_scalar(llvm::IRBuilder<> &B, llvm::Type *fpType,
+                              llvm::Value *ptr) {
+  return B.CreateAlignedLoad(fpType, blas_scalar_ptr(B, fpType, ptr),
+                             blas_scalar_align(B, fpType));
+}
+
+void store_blas_scalar(llvm::IRBuilder<> &B, llvm::Type *fpType,
+                       llvm::Value *ptr, llvm::Value *val) {
+  B.CreateAlignedStore(val, blas_scalar_ptr(B, fpType, ptr),
+                       blas_scalar_align(B, fpType));
+}
+
 SmallVector<llvm::Value *, 1> get_blas_row(llvm::IRBuilder<> &B,
                                            ArrayRef<llvm::Value *> transA,
                                            bool byRef, bool cublas) {

enzyme/Enzyme/Utils.h

@@ -2323,6 +2323,19 @@ void addValueToCache(llvm::Value *arg, bool cache_arg, llvm::Type *ty,
 llvm::Value *load_if_ref(llvm::IRBuilder<> &B, llvm::Type *intType,
                          llvm::Value *V, bool byRef);
 
+/// Complex conjugate of a blas scalar: complex values are represented as
+/// 2-element fp vectors, real values are returned unchanged.
+llvm::Value *complex_conjugate(llvm::IRBuilder<> &B, llvm::Value *V);
+
+/// Load/store a blas scalar of type fpType through a pointer argument, which
+/// may be passed as an integer (julia) or as a pointer to another element
+/// type. Complex scalars are only assumed to be aligned to their component
+/// type, like the C/Fortran complex types they represent.
+llvm::Value *load_blas_scalar(llvm::IRBuilder<> &B, llvm::Type *fpType,
+                              llvm::Value *ptr);
+void store_blas_scalar(llvm::IRBuilder<> &B, llvm::Type *fpType,
+                       llvm::Value *ptr, llvm::Value *val);
+
 void copy_lower_to_upper(llvm::IRBuilder<> &B, llvm::Type *fpType,
                          BlasInfo blas, bool byRef, llvm::Value *layout,
                          llvm::Value *uplo, llvm::Value *A, llvm::Value *lda,

enzyme/test/Enzyme/ForwardMode/blas/cblas_zdotc_sub64_.ll

@@ -0,0 +1,87 @@
+; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s
+
+; Complex conjugated dot product, result returned through the last pointer arg.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare dso_local <2 x double> @__enzyme_fwddiff(...)
+
+declare void @cblas_zdotc_sub64_(i64, <2 x double>*, i64, <2 x double>*, i64, <2 x double>*)
+
+define <2 x double> @active(i64 %len, <2 x double>* noalias %m, <2 x double>* %dm, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, <2 x double>* noalias %m, <2 x double>* %dm, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @inactiveFirst(i64 %len, <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, metadata !"enzyme_const", <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @inactiveSecond(i64 %len, <2 x double>* noalias %m, <2 x double>* noalias %dm, i64 %incm, <2 x double>* noalias %n, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, <2 x double>* noalias %m, <2 x double>* noalias %dm, i64 %incm, metadata !"enzyme_const", <2 x double>* noalias %n, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @f(i64 %len, <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, i64 %incn) {
+entry:
+  %res = alloca <2 x double>, align 8
+  call void @cblas_zdotc_sub64_(i64 %len, <2 x double>* %m, i64 %incm, <2 x double>* %n, i64 %incn, <2 x double>* %res)
+  %r = load <2 x double>, <2 x double>* %res, align 8
+  ret <2 x double> %r
+}
+
+; COM: Pointer types and capture attributes are matched loosely since their
+; COM: spelling depends on the llvm version.
+; CHECK: declare void @cblas_zdotc_sub64_(i64 "enzyme_inactive", [[PTR:(<2 x double>\*|ptr)]] {{(nocapture readonly|readonly captures\(none\))}}, i64 "enzyme_inactive", [[PTR]] {{(nocapture readonly|readonly captures\(none\))}}, i64 "enzyme_inactive", [[PTR]] {{(nocapture|captures\(none\))}})
+
+; CHECK: define <2 x double> @active
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[active:[^(]+]](
+
+; CHECK: define <2 x double> @inactiveFirst
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[inactiveFirst:[^(]+]](
+
+; CHECK: define <2 x double> @inactiveSecond
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[inactiveSecond:[^(]+]](
+
+; COM: dres = dotc(dx, y) + dotc(x, dy), written to the shadow of res (promoted by mem2reg)
+; CHECK: define internal <2 x double> @[[active]](i64 %len, [[PTR]] noalias %m, [[PTR]] %"m'", i64 %incm, [[PTR]] noalias %n, [[PTR]] %"n'", i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotc_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %dotc_sub.ret1 = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %"m'", i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %dotc_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotc_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %"n'", i64 %incn, [[PTR]] %dotc_sub.ret1)
+; CHECK-NEXT:   %1 = load <2 x double>, [[PTR]] %dotc_sub.ret1, align 16
+; CHECK-NEXT:   %2 = fadd fast <2 x double> %0, %1
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %2
+; CHECK-NEXT: }
+
+; CHECK: define internal <2 x double> @[[inactiveFirst]](i64 %len, [[PTR]] noalias %m, i64 %incm, [[PTR]] noalias %n, [[PTR]] %"n'", i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotc_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %"n'", i64 %incn, [[PTR]] %dotc_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotc_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %0
+; CHECK-NEXT: }
+
+; CHECK: define internal <2 x double> @[[inactiveSecond]](i64 %len, [[PTR]] noalias %m, [[PTR]] %"m'", i64 %incm, [[PTR]] noalias %n, i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotc_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %"m'", i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %dotc_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotc_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotc_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %0
+; CHECK-NEXT: }

enzyme/test/Enzyme/ForwardMode/blas/cblas_zdotu_sub64_.ll

@@ -0,0 +1,87 @@
+; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s
+
+; Complex unconjugated dot product, result returned through the last pointer arg.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare dso_local <2 x double> @__enzyme_fwddiff(...)
+
+declare void @cblas_zdotu_sub64_(i64, <2 x double>*, i64, <2 x double>*, i64, <2 x double>*)
+
+define <2 x double> @active(i64 %len, <2 x double>* noalias %m, <2 x double>* %dm, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, <2 x double>* noalias %m, <2 x double>* %dm, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @inactiveFirst(i64 %len, <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, metadata !"enzyme_const", <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, <2 x double>* %dn, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @inactiveSecond(i64 %len, <2 x double>* noalias %m, <2 x double>* noalias %dm, i64 %incm, <2 x double>* noalias %n, i64 %incn) {
+entry:
+  %r = call <2 x double> (...) @__enzyme_fwddiff(<2 x double> (i64, <2 x double>*, i64, <2 x double>*, i64)* @f, i64 %len, <2 x double>* noalias %m, <2 x double>* noalias %dm, i64 %incm, metadata !"enzyme_const", <2 x double>* noalias %n, i64 %incn)
+  ret <2 x double> %r
+}
+
+define <2 x double> @f(i64 %len, <2 x double>* noalias %m, i64 %incm, <2 x double>* noalias %n, i64 %incn) {
+entry:
+  %res = alloca <2 x double>, align 8
+  call void @cblas_zdotu_sub64_(i64 %len, <2 x double>* %m, i64 %incm, <2 x double>* %n, i64 %incn, <2 x double>* %res)
+  %r = load <2 x double>, <2 x double>* %res, align 8
+  ret <2 x double> %r
+}
+
+; COM: Pointer types and capture attributes are matched loosely since their
+; COM: spelling depends on the llvm version.
+; CHECK: declare void @cblas_zdotu_sub64_(i64 "enzyme_inactive", [[PTR:(<2 x double>\*|ptr)]] {{(nocapture readonly|readonly captures\(none\))}}, i64 "enzyme_inactive", [[PTR]] {{(nocapture readonly|readonly captures\(none\))}}, i64 "enzyme_inactive", [[PTR]] {{(nocapture|captures\(none\))}})
+
+; CHECK: define <2 x double> @active
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[active:[^(]+]](
+
+; CHECK: define <2 x double> @inactiveFirst
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[inactiveFirst:[^(]+]](
+
+; CHECK: define <2 x double> @inactiveSecond
+; CHECK-NEXT: entry
+; CHECK-NEXT: call fast <2 x double> @[[inactiveSecond:[^(]+]](
+
+; COM: dres = dotu(dx, y) + dotu(x, dy), written to the shadow of res (promoted by mem2reg)
+; CHECK: define internal <2 x double> @[[active]](i64 %len, [[PTR]] noalias %m, [[PTR]] %"m'", i64 %incm, [[PTR]] noalias %n, [[PTR]] %"n'", i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotu_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %dotu_sub.ret1 = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %"m'", i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %dotu_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotu_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %"n'", i64 %incn, [[PTR]] %dotu_sub.ret1)
+; CHECK-NEXT:   %1 = load <2 x double>, [[PTR]] %dotu_sub.ret1, align 16
+; CHECK-NEXT:   %2 = fadd fast <2 x double> %0, %1
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %2
+; CHECK-NEXT: }
+
+; CHECK: define internal <2 x double> @[[inactiveFirst]](i64 %len, [[PTR]] noalias %m, i64 %incm, [[PTR]] noalias %n, [[PTR]] %"n'", i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotu_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %"n'", i64 %incn, [[PTR]] %dotu_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotu_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %0
+; CHECK-NEXT: }
+
+; CHECK: define internal <2 x double> @[[inactiveSecond]](i64 %len, [[PTR]] noalias %m, [[PTR]] %"m'", i64 %incm, [[PTR]] noalias %n, i64 %incn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %dotu_sub.ret = alloca <2 x double>, align 16
+; CHECK-NEXT:   %res = alloca <2 x double>, align 8
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %"m'", i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %dotu_sub.ret)
+; CHECK-NEXT:   %0 = load <2 x double>, [[PTR]] %dotu_sub.ret, align 16
+; CHECK-NEXT:   call void @cblas_zdotu_sub64_(i64 %len, [[PTR]] %m, i64 %incm, [[PTR]] %n, i64 %incn, [[PTR]] %res)
+; CHECK-NEXT:   ret <2 x double> %0
+; CHECK-NEXT: }