From b11bdaf1156fbc0141de51d68b21fda62ec7db9c Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 13 May 2026 19:39:45 -0700 Subject: [PATCH 1/2] feat: add BLAS triangular solve derivatives --- enzyme/Enzyme/BlasDerivatives.td | 137 +++++++- enzyme/Enzyme/Utils.cpp | 130 ++++++- enzyme/Enzyme/Utils.h | 5 + .../Enzyme/ForwardMode/blas/cblas_trsm_f.ll | 28 ++ .../test/Enzyme/ForwardMode/blas/potrs_f.ll | 45 +++ enzyme/test/Enzyme/ForwardMode/blas/trsm_f.ll | 48 +++ enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll | 42 +++ .../Enzyme/ReverseMode/blas/cblas_trsm_f.ll | 28 ++ enzyme/test/Enzyme/ReverseMode/blas/trsm_f.ll | 49 +++ enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll | 41 +++ enzyme/test/Integration/ForwardMode/blas.cpp | 329 ++++++++++++++++++ enzyme/test/Integration/ReverseMode/blas.cpp | 290 +++++++++++++++ enzyme/test/Integration/blasinfra.h | 97 ++++++ enzyme/tools/enzyme-tblgen/blas-tblgen.cpp | 123 ++++++- enzyme/tools/enzyme-tblgen/blasTAUpdater.h | 6 +- enzyme/tools/enzyme-tblgen/caching.cpp | 5 +- 16 files changed, 1377 insertions(+), 26 deletions(-) create mode 100644 enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll create mode 100644 enzyme/test/Enzyme/ForwardMode/blas/potrs_f.ll create mode 100644 enzyme/test/Enzyme/ForwardMode/blas/trsm_f.ll create mode 100644 enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/blas/cblas_trsm_f.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/blas/trsm_f.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll diff --git a/enzyme/Enzyme/BlasDerivatives.td b/enzyme/Enzyme/BlasDerivatives.td index cfb57caa0ff7..81c4d30aa932 100644 --- a/enzyme/Enzyme/BlasDerivatives.td +++ b/enzyme/Enzyme/BlasDerivatives.td @@ -56,6 +56,7 @@ def Rows : MagicInst; // given a transpose, normal rows, normal cols get the tru def Concat : MagicInst; def ShadowNoInc : MagicInst; +def Dep : MagicInst; class Binop _tys> { string s = _s; @@ -89,8 +90,13 @@ class IntMatchers _before, list; +def side_to_rtrans : IntMatchers< + "charType", "charType", + ["'l'", "'L'", "141", "0", "'R'", "'r'", "142", "1"], + ["'t'", "'T'", "'T'", "'T'", "'N'", "'n'", "'N'", "'N'"] >; def is_diag_int : IntMatchers< @@ -107,6 +113,7 @@ def is_nonunit : MagicInst; def First : MagicInst; def Lookup : MagicInst; def LoadLookup : MagicInst; +def mat_ld : MagicInst; class FAdd { string unused = _tmp; @@ -163,6 +170,10 @@ class MemcpyMatAdd { bit shouldZero = _shouldZero; } +class MatAdd { + bit shouldZero = _shouldZero; +} + class For { string idx = idx_; bit offset = offset_; @@ -672,6 +683,45 @@ def trtrs : CallBlasPattern<(Op $layout, $uplo, $trans, $diag, $n, $nrhs, $A, $l ] >; +def trsv : CallBlasPattern<(Op $layout, $uplo, $trans, $diag, $n, $A, $lda, $x, $incx), + ["x"], + [cblas_layout, uplo, trans, diag, len, mld<["uplo", "n", "n"]>, vinc<["n"]>], + [ + /* A */ + (Seq<["tri", "triangular", "n"], [], 1> + (BlasCall<"lacpy"> $layout, $uplo, $n, $n, (Shadow $A), use<"tri">, $n), + (BlasCall<"ger"> + $layout, + $n, + $n, + Constant<"-1">, + (Rows $trans, + (Concat (Shadow $x)), + (Concat $x)), + (Rows $trans, + (Concat $x), + (Concat (Shadow $x))), + use<"tri">, $n), + + (BlasCall<"copy"> + (ISelect (is_nonunit $diag), ConstantInt<0>, $n), + (First (Shadow $A)), (Add $lda, ConstantInt<1>), + use<"tri">, (Add $n, ConstantInt<1>)), + + (BlasCall<"lacpy"> $layout, $uplo, $n, $n, use<"tri">, $n, (Shadow $A)) + ), + /* x */ (BlasCall<"trsv"> $layout, $uplo, transpose<"trans">, $diag, $n, $A, (ld $A, Char<"N">, $lda, $n, $n), (Shadow $x)), + ] + , + (Seq<["tmp", "vector", "n"], [], 0> + (BlasCall<"copy"> $n, (Dep (Shadow $A), (Dep (Shadow $x), $x)), use<"tmp">, ConstantInt<1>), + (BlasCall<"trmv"> $layout, $uplo, $trans, $diag, $n, (Shadow $A), (Dep (Shadow $x), use<"tmp">), ConstantInt<1>), + (BlasCall<"axpy"> $n, Constant<"-1">, (Dep (Shadow $A), use<"tmp">), ConstantInt<1>, (Shadow $x)), + (BlasCall<"axpy"> (ISelect (is_nonunit $diag), ConstantInt<0>, $n), Constant<"1">, (Dep (Shadow $A), $x), (Shadow $x)), + (BlasCall<"trsv"> $layout, $uplo, $trans, $diag, $n, $A, (ld $A, Char<"N">, $lda, $n, $n), (Shadow $x)) + ) + >; + def spr2 : CallBlasPattern<(Op $layout, $uplo, $n, $alpha, $x, $incx, $y, $incy, $ap), ["ap"], [cblas_layout, uplo, len, fp, vinc<["n"]>, vinc<["n"]>, ap<["n"]>], @@ -718,33 +768,78 @@ def trsm: CallBlasPattern<(Op $layout, $side, $uplo, $transa, $diag, $m, $n, $al [cblas_layout, side, uplo, trans, diag, len, len, fp, mld<["side", "m", "n"]>, mld<["m","n"]>], [ /* alpha */ (AssertingInactiveArg), - /* A */ (AssertingInactiveArg), - /* B */ (AssertingInactiveArg), + /* A */ (Seq<["tmp", "product", "m", "n"], [], 1> + (BlasCall<"lacpy"> $layout, Char<"G">, $m, $n, (Shadow $B), use<"tmp">, (mat_ld $layout, $m, $n)), + (BlasCall<"trsm"> $layout, $side, $uplo, transpose<"transa">, $diag, $m, $n, Constant<"1.0">, $A, (ld $A, (side_to_trans $side), $lda, $n, $m), use<"tmp">, (mat_ld $layout, $m, $n)), + (Seq<["tri", "side_square", "side", "m", "n"], [], 1> + (BlasCall<"lacpy"> $layout, $uplo, (ISelect (is_left $side), $m, $n), (ISelect (is_left $side), $m, $n), (Shadow $A), use<"tri">, (ISelect (is_left $side), $m, $n)), + (BlasCall<"gemm"> + $layout, + (side_to_trans $side), + (side_to_rtrans $side), + (ISelect (is_left $side), $m, $n), + (ISelect (is_left $side), $m, $n), + (ISelect (is_left $side), $n, $m), + Constant<"-1">, + (Rows (side_to_trans $side), + (Rows $transa, + (Concat use<"tmp">, (mat_ld $layout, $m, $n)), + (Concat $B, $ldb)), + (Rows $transa, + (Concat $B, $ldb), + (Concat use<"tmp">, (mat_ld $layout, $m, $n)))), + (Rows (side_to_trans $side), + (Rows $transa, + (Concat $B, $ldb), + (Concat use<"tmp">, (mat_ld $layout, $m, $n))), + (Rows $transa, + (Concat use<"tmp">, (mat_ld $layout, $m, $n)), + (Concat $B, $ldb))), + Constant<"1">, + use<"tri">, (ISelect (is_left $side), $m, $n)), + (BlasCall<"copy"> + (ISelect (is_nonunit $diag), ConstantInt<0>, (ISelect (is_left $side), $m, $n)), + (First (Shadow $A)), (Add $lda, ConstantInt<1>), + use<"tri">, (Add (ISelect (is_left $side), $m, $n), ConstantInt<1>)), + (BlasCall<"lacpy"> $layout, $uplo, (ISelect (is_left $side), $m, $n), (ISelect (is_left $side), $m, $n), use<"tri">, (ISelect (is_left $side), $m, $n), (Shadow $A)) + ) + ), + /* B */ (BlasCall<"trsm"> $layout, $side, $uplo, transpose<"transa">, $diag, $m, $n, $alpha, $A, (ld $A, (side_to_trans $side), $lda, $n, $m), (Shadow $B)), ] + , + (Seq<["tmp", "product", "m", "n"], [], 0> + (Dep (Shadow $alpha), (AssertingInactiveArg)), + (BlasCall<"lascl"> $layout, Char<"G">, ConstantInt<0>, ConstantInt<0>, Constant<"1.0">, $alpha, $m, $n, (Shadow $B), Alloca<1>), + (BlasCall<"lacpy"> $layout, Char<"G">, $m, $n, (Dep (Shadow $A), (Dep (Shadow $B), $B)), $ldb, use<"tmp">, (mat_ld $layout, $m, $n)), + (BlasCall<"trmm"> $layout, $side, $uplo, $transa, $diag, $m, $n, Constant<"1.0">, (Shadow $A), (Dep (Shadow $B), use<"tmp">), (mat_ld $layout, $m, $n)), + (BlasCall<"lascl"> $layout, Char<"G">, ConstantInt<0>, ConstantInt<0>, Constant<"1.0">, Constant<"-1.0">, $m, $n, (Dep (Shadow $A), use<"tmp">), (mat_ld $layout, $m, $n), Alloca<1>), + (MatAdd $layout, $m, $n, (Shadow $B), (Concat (Dep (Shadow $A), use<"tmp">), (mat_ld $layout, $m, $n))), + (MatAdd $layout, (ISelect (is_nonunit $diag), ConstantInt<0>, $m), $n, (Shadow $B), (Concat (Dep (Shadow $A), $B), $ldb)), + (BlasCall<"trsm"> $layout, $side, $uplo, $transa, $diag, $m, $n, Constant<"1.0">, $A, (ld $A, (side_to_trans $side), $lda, $n, $m), (Shadow $B)) + ) >; def uplo_to_normal : IntMatchers< "charType", "charType", - ["'l'", "'L'", "'U'", "'u'"], - ["'n'", "'N'", "'T'", "'t'"] + ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], + ["'n'", "'N'", "'N'", "'N'", "'T'", "'t'", "'T'", "'T'"] >; def uplo_to_trans : IntMatchers< "charType", "charType", - ["'l'", "'L'", "'U'", "'u'"], - ["'t'", "'T'", "'N'", "'n'"] + ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], + ["'t'", "'T'", "'T'", "'T'", "'N'", "'n'", "'N'", "'N'"] >; def uplo_to_side : IntMatchers< "charType", "charType", - ["'l'", "'L'", "'U'", "'u'"], - ["'L'", "'L'", "'R'", "'R'"] + ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], + ["'L'", "'L'", "'L'", "'L'", "'R'", "'R'", "'R'", "'R'"] >; def uplo_to_rside : IntMatchers< "charType", "charType", - ["'l'", "'L'", "'U'", "'u'"], - ["'R'", "'R'", "'L'", "'L'"] + ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], + ["'R'", "'R'", "'R'", "'R'", "'L'", "'L'", "'L'", "'L'"] >; - def potrf: CallBlasPattern<(Op $layout, $uplo, $n, $A, $lda, $info), ["A"], [cblas_layout, uplo, len, mld<["uplo", "n", "n"]>, info], @@ -888,4 +983,20 @@ def potrs: CallBlasPattern<(Op $layout, $uplo, $n, $nrhs, $A, $lda, $B, $ldb, $i ), (BlasCall<"potrs"> $layout, $uplo, $n, $nrhs, $A, (ld $A, Char<"N">, $lda, $n, $n), (Shadow $B), Alloca<1>) ] + , + (Seq<["tmp", "product", "n", "nrhs"], [], 0> + (BlasCall<"lacpy"> $layout, Char<"G">, $n, $nrhs, (Dep (Shadow $A), (Dep (Shadow $B), $B)), $ldb, use<"tmp">, (mat_ld $layout, $n, $nrhs)), + (BlasCall<"trmm"> $layout, Char<"L">, $uplo, (uplo_to_trans $uplo), Char<"N">, $n, $nrhs, Constant<"1.0">, (Dep (Shadow $A), $A), (ld $A, Char<"N">, $lda, $n, $n), (Dep (Shadow $B), use<"tmp">), (mat_ld $layout, $n, $nrhs)), + (BlasCall<"trmm"> $layout, Char<"L">, $uplo, (uplo_to_normal $uplo), Char<"N">, $n, $nrhs, Constant<"1.0">, (Shadow $A), (Dep (Shadow $B), use<"tmp">), (mat_ld $layout, $n, $nrhs)), + (BlasCall<"lascl"> $layout, Char<"G">, ConstantInt<0>, ConstantInt<0>, Constant<"1.0">, Constant<"-1.0">, $n, $nrhs, (Dep (Shadow $A), (Dep (Shadow $B), use<"tmp">)), (mat_ld $layout, $n, $nrhs), Alloca<1>), + (MatAdd $layout, $n, $nrhs, (Shadow $B), (Concat (Dep (Shadow $A), use<"tmp">), (mat_ld $layout, $n, $nrhs))), + + (BlasCall<"lacpy"> $layout, Char<"G">, $n, $nrhs, (Dep (Shadow $A), (Dep (Shadow $B), $B)), $ldb, use<"tmp">, (mat_ld $layout, $n, $nrhs)), + (BlasCall<"trmm"> $layout, Char<"L">, $uplo, (uplo_to_trans $uplo), Char<"N">, $n, $nrhs, Constant<"1.0">, (Shadow $A), (Dep (Shadow $B), use<"tmp">), (mat_ld $layout, $n, $nrhs)), + (BlasCall<"trmm"> $layout, Char<"L">, $uplo, (uplo_to_normal $uplo), Char<"N">, $n, $nrhs, Constant<"1.0">, (Dep (Shadow $A), $A), (ld $A, Char<"N">, $lda, $n, $n), (Dep (Shadow $B), use<"tmp">), (mat_ld $layout, $n, $nrhs)), + (BlasCall<"lascl"> $layout, Char<"G">, ConstantInt<0>, ConstantInt<0>, Constant<"1.0">, Constant<"-1.0">, $n, $nrhs, (Dep (Shadow $A), (Dep (Shadow $B), use<"tmp">)), (mat_ld $layout, $n, $nrhs), Alloca<1>), + (MatAdd $layout, $n, $nrhs, (Shadow $B), (Concat (Dep (Shadow $A), use<"tmp">), (mat_ld $layout, $n, $nrhs))), + + (BlasCall<"potrs"> $layout, $uplo, $n, $nrhs, $A, (ld $A, Char<"N">, $lda, $n, $n), (Shadow $B), Alloca<1>) + ) >; diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp index 6ffb42c296d0..2f64497c6ecb 100644 --- a/enzyme/Enzyme/Utils.cpp +++ b/enzyme/Enzyme/Utils.cpp @@ -2068,6 +2068,129 @@ Function *getOrInsertDifferentialFloatMemcpyMat( return F; } +Function *getOrInsertDifferentialFloatMemcpyMatLayout( + Module &Mod, Type *elementType, PointerType *PT, IntegerType *IT, + IntegerType *CT, unsigned dstalign, unsigned srcalign, bool zeroSrc) { + assert(elementType->isFPOrFPVectorTy()); +#if LLVM_VERSION_MAJOR < 17 +#if LLVM_VERSION_MAJOR >= 15 + if (Mod.getContext().supportsTypedPointers()) { +#endif +#if LLVM_VERSION_MAJOR >= 13 + if (!PT->isOpaquePointerTy()) +#endif + assert(PT->getPointerElementType() == elementType); +#if LLVM_VERSION_MAJOR >= 15 + } +#endif +#endif + std::string name = "__enzyme_dmemcpy_" + tofltstr(elementType) + + "_mat_layout_" + + std::to_string(cast(CT)->getBitWidth()) + + "_" + + std::to_string(cast(IT)->getBitWidth()) + + (zeroSrc ? "_zero" : ""); + FunctionType *FT = FunctionType::get( + Type::getVoidTy(Mod.getContext()), {CT, IT, IT, PT, IT, PT, IT}, false); + + Function *F = cast(Mod.getOrInsertFunction(name, FT).getCallee()); + + if (!F->empty()) + return F; + + F->setLinkage(Function::LinkageTypes::InternalLinkage); +#if LLVM_VERSION_MAJOR >= 16 + F->setOnlyAccessesArgMemory(); +#else + F->addFnAttr(Attribute::ArgMemOnly); +#endif + F->addFnAttr(Attribute::NoUnwind); + F->addFnAttr(Attribute::AlwaysInline); + F->addParamAttr(3, Attribute::NoAlias); + F->addParamAttr(5, Attribute::NoAlias); + + BasicBlock *entry = BasicBlock::Create(F->getContext(), "entry", F); + BasicBlock *init = BasicBlock::Create(F->getContext(), "init.idx", F); + BasicBlock *body = BasicBlock::Create(F->getContext(), "for.body", F); + BasicBlock *initend = BasicBlock::Create(F->getContext(), "init.end", F); + BasicBlock *end = BasicBlock::Create(F->getContext(), "for.end", F); + + auto layout = F->arg_begin(); + layout->setName("layout"); + auto M = layout + 1; + M->setName("M"); + auto N = M + 1; + N->setName("N"); + auto dst = N + 1; + dst->setName("dst"); + auto ldst = dst + 1; + ldst->setName("ldst"); + auto src = ldst + 1; + src->setName("src"); + auto lsrc = src + 1; + lsrc->setName("lsrc"); + + { + IRBuilder<> B(entry); + Value *l0 = B.CreateICmpEQ(M, ConstantInt::get(IT, 0)); + Value *l1 = B.CreateICmpEQ(N, ConstantInt::get(IT, 0)); + B.CreateCondBr(B.CreateOr(l0, l1), end, init); + } + + PHINode *j; + { + IRBuilder<> B(init); + j = B.CreatePHI(IT, 2, "j"); + j->addIncoming(ConstantInt::get(IT, 0), entry); + B.CreateBr(body); + } + + { + IRBuilder<> B(body); + PHINode *i = B.CreatePHI(IT, 2, "i"); + i->addIncoming(ConstantInt::get(IT, 0), init); + + Value *srci = lookup_with_layout(B, elementType, layout, src, lsrc, i, j); + Value *dsti = lookup_with_layout(B, elementType, layout, dst, ldst, i, j); + LoadInst *srcl = B.CreateLoad(elementType, srci, "src.i.l"); + LoadInst *dstl = B.CreateLoad(elementType, dsti, "dst.i.l"); + auto res = B.CreateFAdd(srcl, dstl); + StoreInst *dsts = B.CreateStore(res, dsti); + StoreInst *srcs = nullptr; + if (zeroSrc) + srcs = B.CreateStore(Constant::getNullValue(res->getType()), srci); + if (dstalign) { + dsts->setAlignment(Align(dstalign)); + dstl->setAlignment(Align(dstalign)); + } + if (srcalign) { + if (zeroSrc) + srcs->setAlignment(Align(srcalign)); + srcl->setAlignment(Align(srcalign)); + } + + Value *nexti = + B.CreateAdd(i, ConstantInt::get(IT, 1), "i.next", true, true); + i->addIncoming(nexti, body); + B.CreateCondBr(B.CreateICmpEQ(nexti, M), initend, body); + } + + { + IRBuilder<> B(initend); + Value *nextj = + B.CreateAdd(j, ConstantInt::get(IT, 1), "j.next", true, true); + j->addIncoming(nextj, initend); + B.CreateCondBr(B.CreateICmpEQ(nextj, N), end, init); + } + + { + IRBuilder<> B(end); + B.CreateRetVoid(); + } + + return F; +} + // TODO implement differential memmove Function * getOrInsertDifferentialFloatMemmove(Module &M, Type *T, unsigned dstalign, @@ -3567,6 +3690,7 @@ llvm::Optional extractBLAS(llvm::StringRef in) "dot", "scal", "axpy", "gemv", "gemm", "spmv", "syrk", "nrm2", "trmm", "trmv", "symm", "potrf", "potrs", "copy", "spmv", "syr2k", "potrs", "getrf", "getrs", "trtrs", "getri", "symv", "lacpy", "trsv", + "trsm", }; const char *floatType[] = {"s", "d", "c", "z"}; const char *prefixes[] = {"" /*Fortran*/, "cblas_"}; @@ -4100,7 +4224,11 @@ SmallVector get_blas_row(llvm::IRBuilder<> &B, if (!cublas) { if (!byRef) { - cond = B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 111)); + auto isCblasN = + B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 111)); + auto isN = B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 'N')); + auto isn = B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 'n')); + cond = B.CreateOr(isCblasN, B.CreateOr(isN, isn)); } else { auto isn = B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 'n')); auto isN = B.CreateICmpEQ(trans, ConstantInt::get(trans->getType(), 'N')); diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index c964d52c9427..83c9eccf5776 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -811,6 +811,11 @@ llvm::Function *getOrInsertDifferentialFloatMemcpyMat( llvm::IntegerType *IT, llvm::IntegerType *CT, unsigned dstalign, unsigned srcalign, bool zeroSrc); +llvm::Function *getOrInsertDifferentialFloatMemcpyMatLayout( + llvm::Module &M, llvm::Type *elementType, llvm::PointerType *PT, + llvm::IntegerType *IT, llvm::IntegerType *CT, unsigned dstalign, + unsigned srcalign, bool zeroSrc); + /// Create function for type that performs the derivative memmove on floating /// point memory llvm::Function *getOrInsertDifferentialFloatMemmove( diff --git a/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll b/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll new file mode 100644 index 000000000000..fd2ce3a3ef0c --- /dev/null +++ b/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll @@ -0,0 +1,28 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare dso_local void @__enzyme_fwddiff(...) + +declare void @cblas_dtrsm(i32, i32, i32, i32, i32, i32, i32, double, double*, i32, double*, i32) + +define void @f(i32 %layout, double* %A, double* %B) { +entry: + call void @cblas_dtrsm(i32 %layout, i32 141, i32 121, i32 111, i32 131, i32 4, i32 3, double 2.000000e+00, double* %A, i32 4, double* %B, i32 4) + ret void +} + +define void @active(i32 %layout, double* %A, double* %dA, double* %B, double* %dB) { +entry: + call void (...) @__enzyme_fwddiff(void (i32, double*, double*)* @f, metadata !"enzyme_const", i32 %layout, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %B, double* %dB) + ret void +} + +; CHECK-LABEL: define internal void @fwddiffef( +; CHECK: call void @cblas_dtrsm +; CHECK: icmp eq i32 %layout, 101 +; CHECK: __enzyme_dmemcpy_double_mat_layout_32_32_zero +; CHECK: call void @cblas_dtrsm +; CHECK-NOT: call void bitcast +; CHECK: define internal void @__enzyme_dmemcpy_double_mat_layout_32_32_zero(i32 %layout, diff --git a/enzyme/test/Enzyme/ForwardMode/blas/potrs_f.ll b/enzyme/test/Enzyme/ForwardMode/blas/potrs_f.ll new file mode 100644 index 000000000000..271f3d547e01 --- /dev/null +++ b/enzyme/test/Enzyme/ForwardMode/blas/potrs_f.ll @@ -0,0 +1,45 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -S -enzyme-detect-readthrow=0 | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare dso_local void @__enzyme_fwddiff(...) + +declare void @dpotrs_64_(i8*, i64*, i64*, double*, i64*, double*, i64*, i64*, i64) + +define void @f(double* %A, double* %B) { +entry: + %uplo = alloca i8, align 1 + %n = alloca i64, align 8 + %nrhs = alloca i64, align 8 + %lda = alloca i64, align 8 + %ldb = alloca i64, align 8 + %info = alloca i64, align 8 + store i8 85, i8* %uplo, align 1 + store i64 4, i64* %n, align 8 + store i64 3, i64* %nrhs, align 8 + store i64 4, i64* %lda, align 8 + store i64 4, i64* %ldb, align 8 + call void @dpotrs_64_(i8* %uplo, i64* %n, i64* %nrhs, double* %A, i64* %lda, double* %B, i64* %ldb, i64* %info, i64 1) + ret void +} + +define void @active(double* %A, double* %dA, double* %B, double* %dB) { +entry: + call void (...) @__enzyme_fwddiff(void (double*, double*)* @f, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %B, double* %dB) + ret void +} + +; CHECK-LABEL: define internal void @fwddiffef( +; CHECK: call void @dpotrs_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dtrmm_64_ +; CHECK: call void @dtrmm_64_ +; CHECK: call void @dlascl_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dtrmm_64_ +; CHECK: call void @dtrmm_64_ +; CHECK: call void @dlascl_64_ +; CHECK: call void @dpotrs_64_ +; CHECK: ret void diff --git a/enzyme/test/Enzyme/ForwardMode/blas/trsm_f.ll b/enzyme/test/Enzyme/ForwardMode/blas/trsm_f.ll new file mode 100644 index 000000000000..72fd009c46dc --- /dev/null +++ b/enzyme/test/Enzyme/ForwardMode/blas/trsm_f.ll @@ -0,0 +1,48 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -S -enzyme-detect-readthrow=0 | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare dso_local void @__enzyme_fwddiff(...) + +declare void @dtrsm_64_(i8*, i8*, i8*, i8*, i64*, i64*, double*, double*, i64*, double*, i64*, i64, i64, i64, i64) + +define void @f(double* %A, double* %B) { +entry: + %side = alloca i8, align 1 + %uplo = alloca i8, align 1 + %transa = alloca i8, align 1 + %diag = alloca i8, align 1 + %m = alloca i64, align 8 + %n = alloca i64, align 8 + %alpha = alloca double, align 8 + %lda = alloca i64, align 8 + %ldb = alloca i64, align 8 + store i8 76, i8* %side, align 1 + store i8 85, i8* %uplo, align 1 + store i8 78, i8* %transa, align 1 + store i8 85, i8* %diag, align 1 + store i64 4, i64* %m, align 8 + store i64 3, i64* %n, align 8 + store double 2.000000e+00, double* %alpha, align 8 + store i64 4, i64* %lda, align 8 + store i64 4, i64* %ldb, align 8 + call void @dtrsm_64_(i8* %side, i8* %uplo, i8* %transa, i8* %diag, i64* %m, i64* %n, double* %alpha, double* %A, i64* %lda, double* %B, i64* %ldb, i64 1, i64 1, i64 1, i64 1) + ret void +} + +define void @active(double* %A, double* %dA, double* %B, double* %dB) { +entry: + call void (...) @__enzyme_fwddiff(void (double*, double*)* @f, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %B, double* %dB) + ret void +} + +; CHECK-LABEL: define internal void @fwddiffef( +; CHECK: call void @dtrsm_64_ +; CHECK: call void @dlascl_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dtrmm_64_ +; CHECK: call void @dlascl_64_ +; CHECK: call void @dtrsm_64_ +; CHECK: ret void diff --git a/enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll b/enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll new file mode 100644 index 000000000000..b22fba3cea41 --- /dev/null +++ b/enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll @@ -0,0 +1,42 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -S -enzyme-detect-readthrow=0 | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare dso_local void @__enzyme_fwddiff(...) + +declare void @dtrsv_64_(i8*, i8*, i8*, i64*, double*, i64*, double*, i64*, i64, i64, i64) + +define void @f(double* %A, double* %x) { +entry: + %uplo = alloca i8, align 1 + %trans = alloca i8, align 1 + %diag = alloca i8, align 1 + %n = alloca i64, align 8 + %lda = alloca i64, align 8 + %incx = alloca i64, align 8 + store i8 85, i8* %uplo, align 1 + store i8 78, i8* %trans, align 1 + store i8 85, i8* %diag, align 1 + store i64 4, i64* %n, align 8 + store i64 4, i64* %lda, align 8 + store i64 1, i64* %incx, align 8 + call void @dtrsv_64_(i8* %uplo, i8* %trans, i8* %diag, i64* %n, double* %A, i64* %lda, double* %x, i64* %incx, i64 1, i64 1, i64 1) + ret void +} + +define void @active(double* %A, double* %dA, double* %x, double* %dx) { +entry: + call void (...) @__enzyme_fwddiff(void (double*, double*)* @f, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %x, double* %dx) + ret void +} + +; CHECK-LABEL: define internal void @fwddiffef( +; CHECK: call void @dtrsv_64_ +; CHECK: call void @dcopy_64_ +; CHECK: call void @dtrmv_64_ +; CHECK: call void @daxpy_64_ +; CHECK: call void @daxpy_64_ +; CHECK: call void @dtrsv_64_ +; CHECK: ret void diff --git a/enzyme/test/Enzyme/ReverseMode/blas/cblas_trsm_f.ll b/enzyme/test/Enzyme/ReverseMode/blas/cblas_trsm_f.ll new file mode 100644 index 000000000000..ef2a829f6d8c --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/blas/cblas_trsm_f.ll @@ -0,0 +1,28 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-detect-readthrow=0 -S | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare dso_local void @__enzyme_autodiff(...) + +declare void @cblas_dtrsm(i32, i32, i32, i32, i32, i32, i32, double, double*, i32, double*, i32) + +define void @f(i32 %side, double* %A, double* %B) { +entry: + call void @cblas_dtrsm(i32 102, i32 %side, i32 121, i32 111, i32 131, i32 4, i32 3, double 2.000000e+00, double* %A, i32 4, double* %B, i32 4) + ret void +} + +define void @active(i32 %side, double* %A, double* %dA, double* %B, double* %dB) { +entry: + call void (...) @__enzyme_autodiff(void (i32, double*, double*)* @f, metadata !"enzyme_const", i32 %side, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %B, double* %dB) + ret void +} + +; CHECK-LABEL: define internal void @diffef( +; CHECK: icmp eq i32 %side, 142 +; CHECK: select i1 {{.*}}, i8 84 +; CHECK: icmp eq i32 %side, 142 +; CHECK: select i1 {{.*}}, i8 78 +; CHECK: call void @cblas_dgemm +; CHECK: call void @cblas_dtrsm diff --git a/enzyme/test/Enzyme/ReverseMode/blas/trsm_f.ll b/enzyme/test/Enzyme/ReverseMode/blas/trsm_f.ll new file mode 100644 index 000000000000..17e296b096c8 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/blas/trsm_f.ll @@ -0,0 +1,49 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -S -enzyme-detect-readthrow=0 | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare void @dtrsm_64_(i8*, i8*, i8*, i8*, i64*, i64*, double*, double*, i64*, double*, i64*, i64, i64, i64, i64) + +define void @f(double* %A, double* %B) { +entry: + %side = alloca i8, align 1 + %uplo = alloca i8, align 1 + %transa = alloca i8, align 1 + %diag = alloca i8, align 1 + %m = alloca i64, align 8 + %n = alloca i64, align 8 + %alpha = alloca double, align 8 + %lda = alloca i64, align 8 + %ldb = alloca i64, align 8 + store i8 76, i8* %side, align 1 + store i8 85, i8* %uplo, align 1 + store i8 78, i8* %transa, align 1 + store i8 78, i8* %diag, align 1 + store i64 4, i64* %m, align 8 + store i64 3, i64* %n, align 8 + store double 2.000000e+00, double* %alpha, align 8 + store i64 4, i64* %lda, align 8 + store i64 4, i64* %ldb, align 8 + call void @dtrsm_64_(i8* %side, i8* %uplo, i8* %transa, i8* %diag, i64* %m, i64* %n, double* %alpha, double* %A, i64* %lda, double* %B, i64* %ldb, i64 1, i64 1, i64 1, i64 1) + ret void +} + +declare void @__enzyme_autodiff(...) + +define void @active(double* %A, double* %dA, double* %B, double* %dB) { +entry: + call void (...) @__enzyme_autodiff(void (double*, double*)* @f, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %B, double* %dB) + ret void +} + +; CHECK: define internal void @diffef(double* %A, double* %"A'", double* %B, double* %"B'") +; CHECK: call void @dtrsm_64_ +; CHECK: invertentry: +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dtrsm_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dgemm_64_ +; CHECK: call void @dcopy_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dtrsm_64_ diff --git a/enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll b/enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll new file mode 100644 index 000000000000..8969cf8406aa --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll @@ -0,0 +1,41 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -S -enzyme-detect-readthrow=0 | FileCheck %s + +target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare void @dtrsv_64_(i8*, i8*, i8*, i64*, double*, i64*, double*, i64*, i64, i64, i64) + +define void @f(double* %A, double* %x) { +entry: + %uplo = alloca i8, align 1 + %trans = alloca i8, align 1 + %diag = alloca i8, align 1 + %n = alloca i64, align 8 + %lda = alloca i64, align 8 + %incx = alloca i64, align 8 + store i8 85, i8* %uplo, align 1 + store i8 78, i8* %trans, align 1 + store i8 78, i8* %diag, align 1 + store i64 4, i64* %n, align 8 + store i64 4, i64* %lda, align 8 + store i64 1, i64* %incx, align 8 + call void @dtrsv_64_(i8* %uplo, i8* %trans, i8* %diag, i64* %n, double* %A, i64* %lda, double* %x, i64* %incx, i64 1, i64 1, i64 1) + ret void +} + +declare void @__enzyme_autodiff(...) + +define void @active(double* %A, double* %dA, double* %x, double* %dx) { +entry: + call void (...) @__enzyme_autodiff(void (double*, double*)* @f, metadata !"enzyme_dup", double* %A, double* %dA, metadata !"enzyme_dup", double* %x, double* %dx) + ret void +} + +; CHECK: define internal void @diffef(double* %A, double* %"A'", double* %x, double* %"x'") +; CHECK: call void @dtrsv_64_ +; CHECK: invertentry: +; CHECK: call void @dtrsv_64_ +; CHECK: call void @dlacpy_64_ +; CHECK: call void @dger_64_ +; CHECK: call void @dcopy_64_ +; CHECK: call void @dlacpy_64_ diff --git a/enzyme/test/Integration/ForwardMode/blas.cpp b/enzyme/test/Integration/ForwardMode/blas.cpp index 65bc46aae00e..f8b342320d5d 100644 --- a/enzyme/test/Integration/ForwardMode/blas.cpp +++ b/enzyme/test/Integration/ForwardMode/blas.cpp @@ -35,6 +35,18 @@ void my_dsymv(char layout, char uplo, int N, double alpha, cblas_dsymv(layout, uplo, N, alpha, A, lda, X, incx, beta, Y, incy); } +void my_dtrsv(char layout, char uplo, char trans, char diag, int N, + double *__restrict__ A, int lda, double *__restrict__ X, + int incx) { + cblas_dtrsv(layout, uplo, trans, diag, N, A, lda, X, incx); +} + +void my_dtrsm(char layout, char side, char uplo, char trans, char diag, int M, + int N, double alpha, double *__restrict__ A, int lda, + double *__restrict__ B, int ldb) { + cblas_dtrsm(layout, side, uplo, trans, diag, M, N, alpha, A, lda, B, ldb); +} + double my_ddot(int N, double *__restrict__ X, int incx, double *__restrict__ Y, int incy) { double res = cblas_ddot(N, X, incx, Y, incy); @@ -75,6 +87,13 @@ void my_potrf(char layout, char uplo, int N, double *__restrict__ A, int lda) { cblas_dpotrf(layout, uplo, N, A, lda, nullptr); //&info); } +void my_potrs(char layout, char uplo, int N, int Nrhs, + double *__restrict__ A, int lda, double *__restrict__ B, + int ldb) { + int info; + cblas_dpotrs(layout, uplo, N, Nrhs, A, lda, B, ldb, &info); +} + static void dotTests() { { std::string Test = "DOT active both "; @@ -554,6 +573,218 @@ static void symvTests() { } } +static void trsvTests() { + int N = 17; + for (char layout : {CblasColMajor, CblasRowMajor}) { + for (auto uplo : {'U', 'u', 'L', 'l'}) + for (auto diag : {'U', 'u', 'N', 'n'}) + for (auto transA : + {CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans}) { + BlasInfo inputs[6] = { + /*A*/ BlasInfo(A, layout, N, N, lda), + /*B*/ BlasInfo(B, N, incB), + /*C*/ BlasInfo(), + BlasInfo(), + BlasInfo(), + BlasInfo(), + }; + { + std::string Test = "TRSV active A, x"; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + assert(calls.size() == 1); + assert(calls[0].type == CallType::TRSV); + assert(calls[0].pout_arg1 == B); + assert(calls[0].pin_arg1 == A); + assert(calls[0].layout == layout); + assert(calls[0].targ1 == (char)transA); + assert(calls[0].iarg1 == N); + assert(calls[0].iarg4 == lda); + assert(calls[0].iarg5 == incB); + assert(calls[0].uplo == uplo); + assert(calls[0].diag == diag); + + checkMemoryTrace(inputs, "Primal " + Test, calls); + + init(); + __enzyme_fwddiff( + (void *)my_dtrsv, enzyme_const, layout, enzyme_const, uplo, + enzyme_const, (char)transA, enzyme_const, diag, enzyme_const, + N, enzyme_dup, A, dA, enzyme_const, lda, enzyme_dup, B, dB, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + assert(foundCalls.size() >= 2); + assert(foundCalls[1].type == CallType::COPY); + double *tmp = (double *)foundCalls[1].pout_arg1; + inputs[4] = BlasInfo(tmp, N, 1); + + cblas_dcopy(N, B, incB, tmp, 1); + cblas_dtrmv(layout, uplo, (char)transA, diag, N, dA, lda, tmp, 1); + cblas_daxpy(N, -1.0, tmp, 1, dB, incB); + cblas_daxpy((diag == 'U' || diag == 'u') ? N : 0, 1.0, B, incB, + dB, incB); + cblas_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, dB, incB); + + checkTest(Test); + + SkipVecIncCheck = true; + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + SkipVecIncCheck = false; + } + { + std::string Test = "TRSV active x"; + + init(); + __enzyme_fwddiff( + (void *)my_dtrsv, enzyme_const, layout, enzyme_const, uplo, + enzyme_const, (char)transA, enzyme_const, diag, enzyme_const, + N, enzyme_const, A, enzyme_const, lda, enzyme_dup, B, dB, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + cblas_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, dB, incB); + + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + } + { + std::string Test = "TRSV active A"; + + init(); + __enzyme_fwddiff( + (void *)my_dtrsv, enzyme_const, layout, enzyme_const, uplo, + enzyme_const, (char)transA, enzyme_const, diag, enzyme_const, + N, enzyme_dup, A, dA, enzyme_const, lda, enzyme_const, B, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + } + } + } +} + +static void trsmTests() { + int M = 5; + int N = 4; + for (char layout : {CblasColMajor, CblasRowMajor}) + for (auto side : {'L', 'R'}) + for (auto uplo : {'U', 'L'}) + for (auto diag : {'U', 'N'}) + for (auto transA : + {CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans}) { + int aN = is_left(side) ? M : N; + int tmpLD = mat_ld(layout, M, N); + int localLDA = aN; + int localLDB = tmpLD; + double Abuf[64], dAbuf[64], Bbuf[64], dBbuf[64]; + for (int i = 0; i < 64; i++) { + Abuf[i] = i + 1; + dAbuf[i] = 0.01 * (i + 1); + Bbuf[i] = 0.02 * (i + 1); + dBbuf[i] = 0.03 * (i + 1); + } + double *pA = Abuf; + double *pdA = dAbuf; + double *pB = Bbuf; + double *pdB = dBbuf; + BlasInfo inputs[6] = { + /*A*/ BlasInfo(pA, layout, aN, aN, localLDA), + /*B*/ BlasInfo(pB, layout, M, N, localLDB), + /*C*/ BlasInfo(), + BlasInfo(), + BlasInfo(), + BlasInfo(), + }; + { + std::string Test = "TRSM active A, B"; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, pA, + localLDA, pB, localLDB); + + assert(calls.size() == 1); + assert(calls[0].type == CallType::TRSM); + assert(calls[0].pout_arg1 == pB); + assert(calls[0].pin_arg1 == pA); + assert(calls[0].layout == layout); + assert(calls[0].side == side); + assert(calls[0].uplo == uplo); + assert(calls[0].diag == diag); + assert(calls[0].targ1 == (char)transA); + assert(calls[0].iarg1 == M); + assert(calls[0].iarg2 == N); + assert(calls[0].iarg4 == localLDA); + assert(calls[0].iarg5 == localLDB); + assert(calls[0].farg1 == alpha); + + init(); + __enzyme_fwddiff( + (void *)my_dtrsm, enzyme_const, layout, enzyme_const, side, + enzyme_const, uplo, enzyme_const, (char)transA, enzyme_const, + diag, enzyme_const, M, enzyme_const, N, enzyme_const, alpha, + enzyme_dup, pA, pdA, enzyme_const, localLDA, enzyme_dup, pB, + pdB, enzyme_const, localLDB); + foundCalls = calls; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, pA, + localLDA, pB, localLDB); + + assert(foundCalls.size() >= 3); + assert(foundCalls[2].type == CallType::LACPY); + double *tmp = (double *)foundCalls[2].pout_arg1; + inputs[3] = BlasInfo(tmp, layout, M, N, tmpLD); + + cblas_dlascl(layout, 'G', 0, 0, 1.0, alpha, M, N, pdB, localLDB, + 0); + cblas_dlacpy(layout, 'G', M, N, pB, localLDB, tmp, tmpLD); + cblas_dtrmm(layout, side, uplo, (char)transA, diag, M, N, 1.0, + pdA, localLDA, tmp, tmpLD); + cblas_dlascl(layout, 'G', 0, 0, 1.0, -1.0, M, N, tmp, tmpLD, + 0); + cblas_dtrsm(layout, side, uplo, (char)transA, diag, M, N, 1.0, + pA, localLDA, pdB, localLDB); + + checkTest(Test); + } + { + std::string Test = "TRSM active B"; + init(); + __enzyme_fwddiff( + (void *)my_dtrsm, enzyme_const, layout, enzyme_const, side, + enzyme_const, uplo, enzyme_const, (char)transA, enzyme_const, + diag, enzyme_const, M, enzyme_const, N, enzyme_const, alpha, + enzyme_const, pA, enzyme_const, localLDA, enzyme_dup, pB, + pdB, enzyme_const, localLDB); + foundCalls = calls; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, pA, + localLDA, pB, localLDB); + cblas_dlascl(layout, 'G', 0, 0, 1.0, alpha, M, N, pdB, localLDB, + 0); + cblas_dtrsm(layout, side, uplo, (char)transA, diag, M, N, 1.0, + pA, localLDA, pdB, localLDB); + + checkTest(Test); + } + } +} + static void gemmTests() { // N means normal matrix, T means transposed for (char layout : {CblasRowMajor, CblasColMajor}) { @@ -955,6 +1186,98 @@ static void potrfTests() { } } +static void potrsTests() { + int N = 5; + int Nrhs = 3; + for (char layout : {CblasColMajor, CblasRowMajor}) { + for (auto uplo : {'U', 'u', 'L', 'l'}) { + int tmpLD = mat_ld(layout, N, Nrhs); + int localLDA = N; + int localLDB = tmpLD; + double Abuf[64], dAbuf[64], Bbuf[64], dBbuf[64]; + for (int i = 0; i < 64; i++) { + Abuf[i] = i + 1; + dAbuf[i] = 0.01 * (i + 1); + Bbuf[i] = 0.02 * (i + 1); + dBbuf[i] = 0.03 * (i + 1); + } + double *pA = Abuf; + double *pdA = dAbuf; + double *pB = Bbuf; + double *pdB = dBbuf; + { + std::string Test = "POTRS active A, B"; + init(); + + my_potrs(layout, uplo, N, Nrhs, pA, localLDA, pB, localLDB); + + assert(calls.size() == 1); + assert(calls[0].type == CallType::POTRS); + assert(calls[0].pout_arg1 == pB); + assert(calls[0].pin_arg1 == pA); + assert(calls[0].layout == layout); + assert(calls[0].uplo == uplo); + assert(calls[0].iarg1 == N); + assert(calls[0].iarg2 == Nrhs); + assert(calls[0].iarg4 == localLDA); + assert(calls[0].iarg5 == localLDB); + + init(); + __enzyme_fwddiff( + (void *)my_potrs, enzyme_const, layout, enzyme_const, uplo, + enzyme_const, N, enzyme_const, Nrhs, enzyme_dup, pA, pdA, + enzyme_const, localLDA, enzyme_dup, pB, pdB, enzyme_const, + localLDB); + foundCalls = calls; + init(); + + assert(foundCalls.size() >= 10); + assert(foundCalls[1].type == CallType::LACPY); + double *tmp = (double *)foundCalls[1].pout_arg1; + assert(foundCalls[5].type == CallType::LACPY); + double *tmp2 = (double *)foundCalls[5].pout_arg1; + + my_potrs(layout, uplo, N, Nrhs, pA, localLDA, pB, localLDB); + cblas_dlacpy(layout, 'G', N, Nrhs, pB, localLDB, tmp, tmpLD); + cblas_dtrmm(layout, 'L', uplo, uplo_to_trans(uplo), 'N', N, Nrhs, + 1.0, pA, localLDA, tmp, tmpLD); + cblas_dtrmm(layout, 'L', uplo, uplo_to_normal(uplo), 'N', N, Nrhs, + 1.0, pdA, localLDA, tmp, tmpLD); + cblas_dlascl(layout, 'G', 0, 0, 1.0, -1.0, N, Nrhs, tmp, tmpLD, 0); + + cblas_dlacpy(layout, 'G', N, Nrhs, pB, localLDB, tmp2, tmpLD); + cblas_dtrmm(layout, 'L', uplo, uplo_to_trans(uplo), 'N', N, Nrhs, + 1.0, pdA, localLDA, tmp2, tmpLD); + cblas_dtrmm(layout, 'L', uplo, uplo_to_normal(uplo), 'N', N, Nrhs, + 1.0, pA, localLDA, tmp2, tmpLD); + cblas_dlascl(layout, 'G', 0, 0, 1.0, -1.0, N, Nrhs, tmp2, tmpLD, 0); + cblas_dpotrs(layout, uplo, N, Nrhs, pA, localLDA, pdB, localLDB, + nullptr); + + checkTest(Test); + } + { + std::string Test = "POTRS active B"; + + init(); + __enzyme_fwddiff( + (void *)my_potrs, enzyme_const, layout, enzyme_const, uplo, + enzyme_const, N, enzyme_const, Nrhs, enzyme_const, pA, + enzyme_const, localLDA, enzyme_dup, pB, pdB, enzyme_const, + localLDB); + foundCalls = calls; + init(); + + my_potrs(layout, uplo, N, Nrhs, pA, localLDA, pB, localLDB); + cblas_dpotrs(layout, uplo, N, Nrhs, pA, localLDA, pdB, localLDB, + nullptr); + + checkTest(Test); + } + } + } +} + int main() { dotTests(); @@ -968,5 +1291,11 @@ int main() { potrfTests(); + potrsTests(); + symvTests(); + + trsvTests(); + + trsmTests(); } diff --git a/enzyme/test/Integration/ReverseMode/blas.cpp b/enzyme/test/Integration/ReverseMode/blas.cpp index 6ddbdb7a5fdd..f67a114d74ff 100644 --- a/enzyme/test/Integration/ReverseMode/blas.cpp +++ b/enzyme/test/Integration/ReverseMode/blas.cpp @@ -72,6 +72,13 @@ void my_dtrmv(char layout, char uplo, char trans, inDerivative = true; } +void my_dtrsv(char layout, char uplo, char trans, char diag, int N, + double *__restrict__ A, int lda, double *__restrict__ X, + int incx) { + cblas_dtrsv(layout, uplo, trans, diag, N, A, lda, X, incx); + inDerivative = true; +} + void my_dtrmm(char layout, char side, char uplo, char trans, char diag, int M, int N, double alpha, double * __restrict__ A, int lda, @@ -80,6 +87,13 @@ void my_dtrmm(char layout, char side, char uplo, inDerivative = true; } +void my_dtrsm(char layout, char side, char uplo, char trans, char diag, int M, + int N, double alpha, double *__restrict__ A, int lda, + double *__restrict__ B, int ldb) { + cblas_dtrsm(layout, side, uplo, trans, diag, M, N, alpha, A, lda, B, ldb); + inDerivative = true; +} + void ow_dtrmm(char layout, char side, char uplo, char trans, char diag, int M, int N, double alpha, double * A, int lda, @@ -948,6 +962,278 @@ static void trmvTests() { REALCOPY = false; } +static void trsvTests() { + int N = 17; + for (char layout : {CblasColMajor, CblasRowMajor}) { + for (auto uplo : {'U', 'u', 'L', 'l'}) + for (auto diag : {'U', 'u', 'N', 'n'}) + for (auto transA : + {CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans}) { + BlasInfo inputs[6] = { + /*A*/ BlasInfo(A, layout, N, N, lda), + /*B*/ BlasInfo(B, N, incB), + /*C*/ BlasInfo(), + BlasInfo(), + BlasInfo(), + BlasInfo(), + }; + { + std::string Test = "TRSV active A, x"; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + assert(calls.size() == 1); + assert(calls[0].inDerivative == false); + assert(calls[0].type == CallType::TRSV); + assert(calls[0].pout_arg1 == B); + assert(calls[0].pin_arg1 == A); + assert(calls[0].pin_arg2 == UNUSED_POINTER); + assert(calls[0].farg1 == UNUSED_DOUBLE); + assert(calls[0].farg2 == UNUSED_DOUBLE); + assert(calls[0].layout == layout); + assert(calls[0].targ1 == (char)transA); + assert(calls[0].targ2 == UNUSED_TRANS); + assert(calls[0].iarg1 == N); + assert(calls[0].iarg2 == UNUSED_INT); + assert(calls[0].iarg3 == UNUSED_INT); + assert(calls[0].iarg4 == lda); + assert(calls[0].iarg5 == incB); + assert(calls[0].iarg6 == UNUSED_INT); + assert(calls[0].uplo == uplo); + assert(calls[0].diag == diag); + + checkMemoryTrace(inputs, "Primal " + Test, calls); + + init(); + __enzyme_autodiff((void *)my_dtrsv, enzyme_const, layout, + enzyme_const, uplo, enzyme_const, (char)transA, + enzyme_const, diag, enzyme_const, N, enzyme_dup, + A, dA, enzyme_const, lda, enzyme_dup, B, dB, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + inDerivative = true; + + size_t idx = 1; + double *x = B; + int xinc = incB; + if (foundCalls[idx].type == CallType::COPY) { + x = (double *)foundCalls[idx].pout_arg1; + xinc = 1; + inputs[4] = BlasInfo(x, N, xinc); + cblas_dcopy(N, B, incB, x, xinc); + idx++; + } + + cblas_dtrsv(layout, uplo, (char)transpose(transA), diag, N, A, lda, + dB, incB); + idx++; + + assert(foundCalls[idx].type == CallType::LACPY); + double *tri = (double *)foundCalls[idx].pout_arg1; + inputs[3] = BlasInfo(tri, layout, N, N, N); + + cblas_dlacpy(layout, uplo, N, N, dA, lda, tri, N); + + cblas_dger(layout, N, N, -1.0, is_normal(transA) ? dB : x, + is_normal(transA) ? incB : xinc, + is_normal(transA) ? x : dB, + is_normal(transA) ? xinc : incB, tri, N); + + cblas_dcopy((diag == 'U' || diag == 'u') ? N : 0, dA, lda + 1, tri, + N + 1); + + cblas_dlacpy(layout, uplo, N, N, tri, N, dA, lda); + + checkTest(Test); + + SkipVecIncCheck = true; + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + SkipVecIncCheck = false; + } + { + std::string Test = "TRSV active x"; + + init(); + __enzyme_autodiff((void *)my_dtrsv, enzyme_const, layout, + enzyme_const, uplo, enzyme_const, (char)transA, + enzyme_const, diag, enzyme_const, N, + enzyme_const, A, enzyme_const, lda, enzyme_dup, + B, dB, enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + inDerivative = true; + + cblas_dtrsv(layout, uplo, (char)transpose(transA), diag, N, A, lda, + dB, incB); + + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + } + { + std::string Test = "TRSV active A"; + + init(); + __enzyme_autodiff((void *)my_dtrsv, enzyme_const, layout, + enzyme_const, uplo, enzyme_const, (char)transA, + enzyme_const, diag, enzyme_const, N, enzyme_dup, + A, dA, enzyme_const, lda, enzyme_const, B, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsv(layout, uplo, (char)transA, diag, N, A, lda, B, incB); + + inDerivative = true; + + checkMemoryTrace(inputs, "Expected " + Test, calls); + checkMemoryTrace(inputs, "Found " + Test, foundCalls); + } + } + } +} + +static void trsmTests() { + int M = 5; + int N = 4; + for (char layout : {CblasColMajor, CblasRowMajor}) + for (auto side : {'L', 'R'}) + for (auto uplo : {'U', 'L'}) + for (auto diag : {'U', 'N'}) + for (auto transA : + {CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans}) { + int aN = is_left(side) ? M : N; + int tmpLD = mat_ld(layout, M, N); + BlasInfo inputs[6] = { + /*A*/ BlasInfo(A, layout, aN, aN, lda), + /*B*/ BlasInfo(B, layout, M, N, incB), + /*C*/ BlasInfo(), + BlasInfo(), + BlasInfo(), + BlasInfo(), + }; + { + std::string Test = "TRSM active A, B"; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, A, + lda, B, incB); + + assert(calls.size() == 1); + assert(calls[0].inDerivative == false); + assert(calls[0].type == CallType::TRSM); + assert(calls[0].pout_arg1 == B); + assert(calls[0].pin_arg1 == A); + assert(calls[0].layout == layout); + assert(calls[0].side == side); + assert(calls[0].uplo == uplo); + assert(calls[0].diag == diag); + assert(calls[0].targ1 == (char)transA); + assert(calls[0].iarg1 == M); + assert(calls[0].iarg2 == N); + assert(calls[0].iarg4 == lda); + assert(calls[0].iarg5 == incB); + assert(calls[0].farg1 == alpha); + + init(); + __enzyme_autodiff( + (void *)my_dtrsm, enzyme_const, layout, enzyme_const, side, + enzyme_const, uplo, enzyme_const, (char)transA, enzyme_const, + diag, enzyme_const, M, enzyme_const, N, enzyme_const, alpha, + enzyme_dup, A, dA, enzyme_const, lda, enzyme_dup, B, dB, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, A, + lda, B, incB); + + size_t idx = 1; + double *B0 = B; + int B0LD = incB; + if (foundCalls[idx].type == CallType::LACPY && + foundCalls[idx].pin_arg1 == B) { + B0 = (double *)foundCalls[idx].pout_arg1; + B0LD = M; + inputs[3] = BlasInfo(B0, layout, M, N, B0LD); + cblas_dlacpy(layout, '\0', M, N, B, incB, B0, B0LD); + idx++; + } + + inDerivative = true; + + assert(foundCalls[idx].type == CallType::LACPY); + double *tmp = (double *)foundCalls[idx].pout_arg1; + inputs[4] = BlasInfo(tmp, layout, M, N, tmpLD); + cblas_dlacpy(layout, 'G', M, N, dB, incB, tmp, tmpLD); + idx++; + + cblas_dtrsm(layout, side, uplo, (char)transpose(transA), diag, M, + N, 1.0, A, lda, tmp, tmpLD); + idx++; + + assert(foundCalls[idx].type == CallType::LACPY); + int triN = is_left(side) ? M : N; + double *tri = (double *)foundCalls[idx].pout_arg1; + inputs[5] = BlasInfo(tri, layout, triN, triN, triN); + cblas_dlacpy(layout, uplo, triN, triN, dA, lda, tri, triN); + + bool normal = is_normal(transA); + double *gemmA = + is_left(side) ? (normal ? tmp : B0) : (normal ? B0 : tmp); + int gemmALD = + is_left(side) ? (normal ? tmpLD : B0LD) + : (normal ? B0LD : tmpLD); + double *gemmB = + is_left(side) ? (normal ? B0 : tmp) : (normal ? tmp : B0); + int gemmBLD = + is_left(side) ? (normal ? B0LD : tmpLD) + : (normal ? tmpLD : B0LD); + cblas_dgemm(layout, side_to_trans(side), side_to_rtrans(side), + triN, triN, is_left(side) ? N : M, -1.0, gemmA, + gemmALD, gemmB, gemmBLD, 1.0, tri, triN); + + cblas_dcopy((diag == 'U') ? triN : 0, dA, lda + 1, tri, + triN + 1); + cblas_dlacpy(layout, uplo, triN, triN, tri, triN, dA, lda); + cblas_dtrsm(layout, side, uplo, (char)transpose(transA), diag, M, + N, alpha, A, lda, dB, incB); + + checkTest(Test); + } + { + std::string Test = "TRSM active B"; + + init(); + __enzyme_autodiff( + (void *)my_dtrsm, enzyme_const, layout, enzyme_const, side, + enzyme_const, uplo, enzyme_const, (char)transA, enzyme_const, + diag, enzyme_const, M, enzyme_const, N, enzyme_const, alpha, + enzyme_const, A, enzyme_const, lda, enzyme_dup, B, dB, + enzyme_const, incB); + foundCalls = calls; + init(); + + my_dtrsm(layout, side, uplo, (char)transA, diag, M, N, alpha, A, + lda, B, incB); + + inDerivative = true; + cblas_dtrsm(layout, side, uplo, (char)transpose(transA), diag, M, + N, alpha, A, lda, dB, incB); + + checkTest(Test); + } + } +} + static void trmmTests() { // N means normal matrix, T means transposed // TODO: row major is presently an exepcted failure. We should re-enable. @@ -2215,6 +2501,10 @@ int main() { trmvTests(); + trsvTests(); + + trsmTests(); + trmmTests(); syrkTests(); diff --git a/enzyme/test/Integration/blasinfra.h b/enzyme/test/Integration/blasinfra.h index 19fa5611cd15..2f5064c9a451 100644 --- a/enzyme/test/Integration/blasinfra.h +++ b/enzyme/test/Integration/blasinfra.h @@ -106,6 +106,26 @@ char side_to_trans(char c) { } } +char side_to_rtrans(char c) { + switch (c) { + case 'L': + return 'T'; + case 'l': + return 't'; + case 'R': + return 'N'; + case 'r': + return 'n'; + default: + printf("Illegal side_to_rtrans of '%c' %d\n", c, c); + exit(1); + } +} + +int mat_ld(char layout, int rows, int cols) { + return layout == 101 ? cols : rows; +} + bool is_normal(char c) { switch (c) { case 'N': @@ -344,6 +364,7 @@ enum class CallType { LACPY, MEMSET, TRMV, + TRSV, TRMM, SYRK, SYR2K, @@ -481,6 +502,9 @@ void printty(CallType v) { case CallType::TRMV: printf("TRMV"); return; + case CallType::TRSV: + printf("TRSV"); + return; case CallType::TRMM: printf("TRMM"); return; @@ -982,6 +1006,31 @@ void printcall(BlasCall rcall) { printty(rcall.iarg5); printf(")"); return; + case CallType::TRSV: + printf("TRSV(abi="); + printty(rcall.abi); + printf(", handle="); + printty(rcall.handle); + printf(", layout="); + printty(rcall.layout); + printf(", uplo="); + printty(rcall.uplo); + printf(", trans="); + printty(rcall.targ1); + printf(", diag="); + printty(rcall.diag); + printf(", N="); + printty(rcall.iarg1); + printf(", A="); + printty(rcall.pin_arg1); + printf(", lda="); + printty(rcall.iarg4); + printf(", X="); + printty(rcall.pout_arg1); + printf(", incx="); + printty(rcall.iarg5); + printf(")"); + return; case CallType::TRMM: printf("TRMM(abi="); printty(rcall.abi); @@ -2027,6 +2076,34 @@ __attribute__((noinline)) void cblas_dtrmv(char layout, char uplo, char trans, calls.push_back(call); } +__attribute__((noinline)) void cblas_dtrsv(char layout, char uplo, char trans, + char diag, int N, double *A, int lda, + double *X, int incx) { + BlasCall call = {ABIType::CBLAS, + UNUSED_HANDLE, + inDerivative, + CallType::TRSV, + X, + A, + UNUSED_POINTER, + UNUSED_DOUBLE, + UNUSED_DOUBLE, + layout, + trans, + UNUSED_TRANS, + N, + UNUSED_INT, + UNUSED_INT, + lda, + incx, + UNUSED_INT, + UNUSED_INT, + UNUSED_TRANS, + uplo, + diag}; + calls.push_back(call); +} + // B := alpha*op( A )*B, or B := alpha*B*op( A ), __attribute__((noinline)) void cblas_dtrmm(char layout, char side, char uplo, char trans, char diag, int M, int N, @@ -2873,6 +2950,26 @@ void checkMemory(BlasCall rcall, BlasInfo inputs[6], std::string test, uplo_char == 'l'); return; } + case CallType::TRSV: { + auto X = pointer_to_index(rcall.pout_arg1, inputs); + auto A = pointer_to_index(rcall.pin_arg1, inputs); + + auto layout = rcall.layout; + auto diag_char = rcall.diag; + auto uplo_char = rcall.uplo; + auto N = rcall.iarg1; + auto lda = rcall.iarg4; + auto incX = rcall.iarg5; + + checkMatrix(A, "A", layout, /*rows=*/N, /*cols=*/N, /*ld=*/lda, test, + rcall, trace); + checkVector(X, "X", /*len=*/N, /*inc=*/incX, test, rcall, trace); + + checkDiag(diag_char, test, rcall, trace); + assert(uplo_char == 'U' || uplo_char == 'u' || uplo_char == 'L' || + uplo_char == 'l'); + return; + } case CallType::TRMM: { // B := alpha*op( A )*B, or B := alpha*B*op( A ), auto B = pointer_to_index(rcall.pout_arg1, inputs); diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp index dbdcc0f6a6f6..eb68eb89c950 100644 --- a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp @@ -709,7 +709,7 @@ void emit_extract_calls(const TGPattern &pattern, raw_ostream &os) { os << " if (byRefFloat) {\n"; extract_scalar(name, "fpType", os); os << " }\n"; - } else if (ty == ArgType::trans) { + } else if (is_char_arg(ty)) { // we are in the byRef branch and trans only exist in lv23. // So just unconditionally asume that no layout exist and use i-1 os << " if (byRef) {\n"; @@ -951,6 +951,12 @@ void rev_call_arg(bool forward, const DagInit *ruleDag, os << "); })"; return; } + if (Def->getName() == "Dep") { + if (Dag->getNumArgs() != 2) + PrintFatalError(pattern.getLoc(), "only 2-arg Dep operands supported"); + rev_call_arg(forward, Dag, pattern, 1, os, vars); + return; + } if (Def->getName() == "ld") { if (Dag->getNumArgs() != 5) PrintFatalError(pattern.getLoc(), "only 5-arg ld operands supported"); @@ -965,6 +971,31 @@ void rev_call_arg(bool forward, const DagInit *ruleDag, << ", cache_" << matName << ", byRef, cublas)}"; return; } + if (Def->getName() == "mat_ld") { + if (Dag->getNumArgs() != 3) + PrintFatalError(pattern.getLoc(), + "only 3-arg mat_ld operands supported"); + const auto rowName = Dag->getArgNameStr(1); + const auto colName = Dag->getArgNameStr(2); + os << "({"; + os << " auto rows = load_if_ref(Builder2, intType, arg_" << rowName + << ", byRef);\n"; + os << " auto cols = load_if_ref(Builder2, intType, arg_" << colName + << ", byRef);\n"; + os << " Value *res = rows;\n"; + os << " if (cblas) {\n"; + os << " auto layout = load_if_ref(Builder2, charType, arg_layout, " + "byRef);\n"; + os << " auto is_row = Builder2.CreateICmpEQ(layout, " + "ConstantInt::get(layout->getType(), 101));\n"; + os << " res = CreateSelect(Builder2, is_row, cols, rows);\n"; + os << " }\n"; + os << " SmallVector vals = {to_blas_callconv(Builder2, " + "res, byRef, cublas, julia_decl_type, allocationBuilder, " + "\"mat.ld\")};\n"; + os << " vals; })"; + return; + } if (Def->getName() == "is_zero") { if (Dag->getNumArgs() != 1) PrintFatalError(pattern.getLoc(), "only 1-arg ld operands supported"); @@ -1493,7 +1524,7 @@ void rev_call_args(bool forward, Twine argName, const TGPattern &pattern, n = 1; if (func == "gemm" || func == "syrk" || func == "syr2k" || func == "symm") n = 2; - if (func == "trmv" || func == "trtrs") + if (func == "trmv" || func == "trsv" || func == "trtrs") n = 3; if (func == "trmm" || func == "trsm") n = 4; @@ -1544,7 +1575,8 @@ void emit_tmp_creation(const Record *Def, raw_ostream &os, StringRef builder) { auto action = args[1]; assert(action == "product" || action == "is_normal" || action == "triangular" || action == "vector" || - action == "zerotriangular" || action == "is_left"); + action == "zerotriangular" || action == "is_left" || + action == "side_square"); if (action == "product") { const auto matName = args[0]; const auto dim1 = "arg_" + args[2]; @@ -1581,6 +1613,22 @@ void emit_tmp_creation(const Record *Def, raw_ostream &os, StringRef builder) { os << " Value *size_" << vecName << " = " << builder << ".CreateSelect(is_left(" << builder << ", " << trans << ", byRef, cublas), len1, len2);\n"; + } else if (action == "side_square") { + assert(args.size() == 5); + const auto matName = args[0]; + const auto side = "arg_" + args[2]; + const auto dim1 = "arg_" + args[3]; + const auto dim2 = "arg_" + args[4]; + os << " Value *len1 = load_if_ref(" << builder << ", intType," << dim1 + << ", byRef);\n" + << " Value *len2 = load_if_ref(" << builder << ", intType," << dim2 + << ", byRef);\n"; + os << " Value *side_len_" << matName << " = " << builder + << ".CreateSelect(is_left(" << builder << ", " << side + << ", byRef, cublas), len1, len2);\n"; + os << " Value *size_" << matName << " = " << builder + << ".CreateMul(side_len_" << matName << ", side_len_" << matName + << ");\n"; } else if (action == "vector") { assert(args.size() == 3); const auto vecName = args[0]; @@ -1733,6 +1781,18 @@ void emit_dag(bool forward, Twine resultVarName, const DagInit *ruleDag, } return; } + if (Def->getName() == "Dep") { + if (ruleDag->getNumArgs() != 2) + PrintFatalError(pattern.getLoc(), "only 2-arg Dep operands supported"); + if (forward) + emit_if_rule_condition(pattern, ruleDag, "", " ", os); + emit_dag(forward, resultVarName, cast(ruleDag->getArg(1)), + argName + "_dep", os, argName, actArg, pattern, runtimeChecked, + vars); + if (forward) + os << " }\n"; + return; + } if (Def->isSubClassOf("BlasCall")) { if (forward) emit_if_rule_condition(pattern, ruleDag, "", " ", os); @@ -1967,6 +2027,56 @@ void emit_dag(bool forward, Twine resultVarName, const DagInit *ruleDag, os << " }\n"; return; } + if (Def->isSubClassOf("MatAdd")) { + assert(ruleDag->getNumArgs() == 5); + if (forward) + emit_if_rule_condition(pattern, ruleDag, "", " ", os); + auto ty = ArgType::len; + os << " {\n"; + if (!forward && !runtimeChecked) + emit_runtime_condition(ruleDag, argName, " ", "Builder2", + (ty == ArgType::fp), os); + os << " // MatAdd\n"; + for (size_t i = 1; i < ruleDag->getNumArgs(); i++) { + os << " SmallVector " << argPrefix << "_" << i + << ";\n"; + os << " for (auto v : "; + rev_call_arg(forward, ruleDag, pattern, i, os, vars); + os << ") " << argPrefix << "_" << i << ".push_back(v);\n"; + } + + os << " const auto Defs = gutils->getInvertedBundles(&call, {" + << ValueType_helper(pattern, actArg, ruleDag) + << "}, Builder2, /* lookup */ " << (!forward) << ");\n"; + os << " Value *layout = cblas ? load_if_ref(Builder2, charType, " + "arg_layout, byRef) : ConstantInt::get(charType, 102);\n"; + os << " auto layoutType = cast(layout->getType());\n"; + os << " auto M = load_if_ref(Builder2, intType, " << argPrefix + << "_1[0], byRef);\n"; + os << " auto N = load_if_ref(Builder2, intType, " << argPrefix + << "_2[0], byRef);\n"; + os << " auto dstLD = load_if_ref(Builder2, intType, " << argPrefix + << "_3[1], byRef);\n"; + os << " auto srcLD = load_if_ref(Builder2, intType, " << argPrefix + << "_4[1], byRef);\n"; + os << " auto dmemcpymat = " + "getOrInsertDifferentialFloatMemcpyMatLayout(*gutils->oldFunc->" + "getParent(), fpType, cast(" + << argPrefix << "_3[0]->getType()), intType, layoutType, 0, 0, " + << Def->getValueAsBit("shouldZero") << ");\n"; + os << " Value *matadd_args[] = {layout, M, N, " << argPrefix + << "_3[0], dstLD, " << argPrefix << "_4[0], srcLD};\n"; + os << " auto cubcall = cast(Builder2.CreateCall(" + "dmemcpymat, matadd_args, Defs));\n"; + + if (!forward && !runtimeChecked) + emit_runtime_continue(ruleDag, argName, " ", "Builder2", + (ty == ArgType::fp), os); + os << " }\n"; + if (forward) + os << " }\n"; + return; + } if (Def->getName() == "FirstUse" || Def->isSubClassOf("FirstUse")) { os << " {\n"; os << " // FirstUse\n"; @@ -2354,9 +2464,10 @@ void emit_rev_rewrite_rules(const StringMap &patternMap, os << " auto bb_name = Builder2.GetInsertBlock()->getName();\n"; for (size_t iteri = 0; iteri < activeArgs.size(); iteri++) { - // trtrs do in reversed arg order. - size_t i = (pattern.getName() != "trtrs") ? iteri - : (activeArgs.size() - 1 - iteri); + // TRSV/TRTRS need the solved cotangent before forming dA. + size_t i = (pattern.getName() != "trtrs" && pattern.getName() != "trsv") + ? iteri + : (activeArgs.size() - 1 - iteri); StringRef extraCond; auto rule = rules[i]; const size_t actArg = activeArgs[i]; diff --git a/enzyme/tools/enzyme-tblgen/blasTAUpdater.h b/enzyme/tools/enzyme-tblgen/blasTAUpdater.h index 2ff92d047b06..d48546d678db 100644 --- a/enzyme/tools/enzyme-tblgen/blasTAUpdater.h +++ b/enzyme/tools/enzyme-tblgen/blasTAUpdater.h @@ -154,12 +154,10 @@ inline void emit_BLASTA(TGPattern &pattern, raw_ostream &os) { break; case ArgType::uplo: case ArgType::trans: - os << " updateAnalysis(call.getArgOperand(" << i - << " + offset), ttChar, &call);\n"; - break; case ArgType::diag: case ArgType::side: - // TODO + os << " updateAnalysis(call.getArgOperand(" << i + << " + offset), ttChar, &call);\n"; break; } } diff --git a/enzyme/tools/enzyme-tblgen/caching.cpp b/enzyme/tools/enzyme-tblgen/caching.cpp index 0c0edec1fea4..7a10ccfdc234 100644 --- a/enzyme/tools/enzyme-tblgen/caching.cpp +++ b/enzyme/tools/enzyme-tblgen/caching.cpp @@ -333,7 +333,8 @@ void emit_cache_for_reverse(const TGPattern &pattern, raw_ostream &os) { << " if ((Mode == DerivativeMode::ReverseModeCombined ||\n" << " Mode == DerivativeMode::ReverseModePrimal) && cachetype) {\n" << " SmallVector cacheValues;\n"; -if (pattern.getName() == "potrf" || pattern.getName() == "trtrs") { +if (pattern.getName() == "potrf" || pattern.getName() == "trtrs" || + pattern.getName() == "trsv" || pattern.getName() == "trsm") { os << "BuilderZ.SetInsertPoint(gutils->getNewFromOriginal(&call)->getNextNode());\n"; } os << " if (byRef) {\n"; @@ -344,7 +345,7 @@ os << "BuilderZ.SetInsertPoint(gutils->getNewFromOriginal(&call)->getNextNode()) const char* scalType; if (ty == ArgType::len || ty == ArgType::vincInc || ty == ArgType::mldLD) { scalType = "intType"; - } else if (ty == ArgType::trans) { + } else if (is_char_arg(ty)) { scalType = "charType"; } else if (ty == ArgType::fp) { scalType = "fpType"; From 8f3163ccf1d3fa91ed808660d2cbbb6a00aac42f Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Thu, 14 May 2026 14:01:38 -0700 Subject: [PATCH 2/2] fix: simplify triangular BLAS helper --- enzyme/Enzyme/BlasDerivatives.td | 16 +- enzyme/Enzyme/Utils.cpp | 154 ++---------------- enzyme/Enzyme/Utils.h | 9 +- .../Enzyme/ForwardMode/blas/cblas_trsm_f.ll | 4 +- enzyme/test/Enzyme/ReverseMode/blas/dlacpy.ll | 95 ++--------- enzyme/tools/enzyme-tblgen/blas-tblgen.cpp | 37 +++-- 6 files changed, 65 insertions(+), 250 deletions(-) diff --git a/enzyme/Enzyme/BlasDerivatives.td b/enzyme/Enzyme/BlasDerivatives.td index 81c4d30aa932..aac6c84dc722 100644 --- a/enzyme/Enzyme/BlasDerivatives.td +++ b/enzyme/Enzyme/BlasDerivatives.td @@ -821,23 +821,23 @@ def trsm: CallBlasPattern<(Op $layout, $side, $uplo, $transa, $diag, $m, $n, $al def uplo_to_normal : IntMatchers< "charType", "charType", - ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], - ["'n'", "'N'", "'N'", "'N'", "'T'", "'t'", "'T'", "'T'"] + ["'l'", "'L'", "'U'", "'u'"], + ["'n'", "'N'", "'T'", "'t'"] >; def uplo_to_trans : IntMatchers< "charType", "charType", - ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], - ["'t'", "'T'", "'T'", "'T'", "'N'", "'n'", "'N'", "'N'"] + ["'l'", "'L'", "'U'", "'u'"], + ["'t'", "'T'", "'N'", "'n'"] >; def uplo_to_side : IntMatchers< "charType", "charType", - ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], - ["'L'", "'L'", "'L'", "'L'", "'R'", "'R'", "'R'", "'R'"] + ["'l'", "'L'", "'U'", "'u'"], + ["'L'", "'L'", "'R'", "'R'"] >; def uplo_to_rside : IntMatchers< "charType", "charType", - ["'l'", "'L'", "122", "0", "'U'", "'u'", "121", "1"], - ["'R'", "'R'", "'R'", "'R'", "'L'", "'L'", "'L'", "'L'"] + ["'l'", "'L'", "'U'", "'u'"], + ["'R'", "'R'", "'L'", "'L'"] >; def potrf: CallBlasPattern<(Op $layout, $uplo, $n, $A, $lda, $info), diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp index 2f64497c6ecb..b82ee278f7cd 100644 --- a/enzyme/Enzyme/Utils.cpp +++ b/enzyme/Enzyme/Utils.cpp @@ -1912,7 +1912,8 @@ Function *getOrInsertMemcpyMat(Module &Mod, Type *elementType, PointerType *PT, Function *getOrInsertDifferentialFloatMemcpyMat( Module &Mod, Type *elementType, PointerType *PT, IntegerType *IT, - IntegerType *CT, unsigned dstalign, unsigned srcalign, bool zeroSrc) { + IntegerType *UT, IntegerType *LT, unsigned dstalign, unsigned srcalign, + bool zeroSrc) { assert(elementType->isFPOrFPVectorTy()); #if LLVM_VERSION_MAJOR < 17 #if LLVM_VERSION_MAJOR >= 15 @@ -1927,10 +1928,13 @@ Function *getOrInsertDifferentialFloatMemcpyMat( #endif #endif std::string name = "__enzyme_dmemcpy_" + tofltstr(elementType) + "_mat_" + + std::to_string(cast(UT)->getBitWidth()) + + "_" + std::to_string(cast(LT)->getBitWidth()) + + "_" + std::to_string(cast(IT)->getBitWidth()) + (zeroSrc ? "_zero" : ""); FunctionType *FT = FunctionType::get(Type::getVoidTy(Mod.getContext()), - {CT, IT, IT, PT, IT, PT, IT}, false); + {UT, LT, IT, IT, PT, IT, PT, IT}, false); Function *F = cast(Mod.getOrInsertFunction(name, FT).getCallee()); @@ -1945,8 +1949,8 @@ Function *getOrInsertDifferentialFloatMemcpyMat( #endif F->addFnAttr(Attribute::NoUnwind); F->addFnAttr(Attribute::AlwaysInline); - F->addParamAttr(3, Attribute::NoAlias); - F->addParamAttr(5, Attribute::NoAlias); + F->addParamAttr(4, Attribute::NoAlias); + F->addParamAttr(6, Attribute::NoAlias); BasicBlock *entry = BasicBlock::Create(F->getContext(), "entry", F); BasicBlock *swtch = BasicBlock::Create(F->getContext(), "swtch", F); @@ -1957,7 +1961,9 @@ Function *getOrInsertDifferentialFloatMemcpyMat( auto uplo = F->arg_begin(); uplo->setName("uplo"); - auto M = uplo + 1; + auto layout = uplo + 1; + layout->setName("layout"); + auto M = layout + 1; M->setName("M"); auto N = M + 1; N->setName("N"); @@ -1982,8 +1988,8 @@ Function *getOrInsertDifferentialFloatMemcpyMat( { IRBuilder<> B(swtch); auto swtchT = B.CreateSwitch(uplo, Ginit); - swtchT->addCase(ConstantInt::get(CT, 'U'), Uinit); - swtchT->addCase(ConstantInt::get(CT, 'L'), Linit); + swtchT->addCase(ConstantInt::get(UT, 'U'), Uinit); + swtchT->addCase(ConstantInt::get(UT, 'L'), Linit); } std::pair todo[] = { @@ -2019,15 +2025,8 @@ Function *getOrInsertDifferentialFloatMemcpyMat( PHINode *i = B.CreatePHI(IT, 2, dir + "i"); i->addIncoming(istart, init); - Value *srci = B.CreateInBoundsGEP( - elementType, src, - B.CreateAdd(i, B.CreateMul(j, lsrc, "", true, true), "", true, true), - dir + "src.i"); - - Value *dsti = B.CreateInBoundsGEP( - elementType, dst, - B.CreateAdd(i, B.CreateMul(j, ldst, "", true, true), "", true, true), - dir + "dst.i"); + Value *srci = lookup_with_layout(B, elementType, layout, src, lsrc, i, j); + Value *dsti = lookup_with_layout(B, elementType, layout, dst, ldst, i, j); LoadInst *srcl = B.CreateLoad(elementType, srci, dir + "src.i.l"); LoadInst *dstl = B.CreateLoad(elementType, dsti, dir + "dst.i.l"); auto res = B.CreateFAdd(srcl, dstl); @@ -2068,129 +2067,6 @@ Function *getOrInsertDifferentialFloatMemcpyMat( return F; } -Function *getOrInsertDifferentialFloatMemcpyMatLayout( - Module &Mod, Type *elementType, PointerType *PT, IntegerType *IT, - IntegerType *CT, unsigned dstalign, unsigned srcalign, bool zeroSrc) { - assert(elementType->isFPOrFPVectorTy()); -#if LLVM_VERSION_MAJOR < 17 -#if LLVM_VERSION_MAJOR >= 15 - if (Mod.getContext().supportsTypedPointers()) { -#endif -#if LLVM_VERSION_MAJOR >= 13 - if (!PT->isOpaquePointerTy()) -#endif - assert(PT->getPointerElementType() == elementType); -#if LLVM_VERSION_MAJOR >= 15 - } -#endif -#endif - std::string name = "__enzyme_dmemcpy_" + tofltstr(elementType) + - "_mat_layout_" + - std::to_string(cast(CT)->getBitWidth()) + - "_" + - std::to_string(cast(IT)->getBitWidth()) + - (zeroSrc ? "_zero" : ""); - FunctionType *FT = FunctionType::get( - Type::getVoidTy(Mod.getContext()), {CT, IT, IT, PT, IT, PT, IT}, false); - - Function *F = cast(Mod.getOrInsertFunction(name, FT).getCallee()); - - if (!F->empty()) - return F; - - F->setLinkage(Function::LinkageTypes::InternalLinkage); -#if LLVM_VERSION_MAJOR >= 16 - F->setOnlyAccessesArgMemory(); -#else - F->addFnAttr(Attribute::ArgMemOnly); -#endif - F->addFnAttr(Attribute::NoUnwind); - F->addFnAttr(Attribute::AlwaysInline); - F->addParamAttr(3, Attribute::NoAlias); - F->addParamAttr(5, Attribute::NoAlias); - - BasicBlock *entry = BasicBlock::Create(F->getContext(), "entry", F); - BasicBlock *init = BasicBlock::Create(F->getContext(), "init.idx", F); - BasicBlock *body = BasicBlock::Create(F->getContext(), "for.body", F); - BasicBlock *initend = BasicBlock::Create(F->getContext(), "init.end", F); - BasicBlock *end = BasicBlock::Create(F->getContext(), "for.end", F); - - auto layout = F->arg_begin(); - layout->setName("layout"); - auto M = layout + 1; - M->setName("M"); - auto N = M + 1; - N->setName("N"); - auto dst = N + 1; - dst->setName("dst"); - auto ldst = dst + 1; - ldst->setName("ldst"); - auto src = ldst + 1; - src->setName("src"); - auto lsrc = src + 1; - lsrc->setName("lsrc"); - - { - IRBuilder<> B(entry); - Value *l0 = B.CreateICmpEQ(M, ConstantInt::get(IT, 0)); - Value *l1 = B.CreateICmpEQ(N, ConstantInt::get(IT, 0)); - B.CreateCondBr(B.CreateOr(l0, l1), end, init); - } - - PHINode *j; - { - IRBuilder<> B(init); - j = B.CreatePHI(IT, 2, "j"); - j->addIncoming(ConstantInt::get(IT, 0), entry); - B.CreateBr(body); - } - - { - IRBuilder<> B(body); - PHINode *i = B.CreatePHI(IT, 2, "i"); - i->addIncoming(ConstantInt::get(IT, 0), init); - - Value *srci = lookup_with_layout(B, elementType, layout, src, lsrc, i, j); - Value *dsti = lookup_with_layout(B, elementType, layout, dst, ldst, i, j); - LoadInst *srcl = B.CreateLoad(elementType, srci, "src.i.l"); - LoadInst *dstl = B.CreateLoad(elementType, dsti, "dst.i.l"); - auto res = B.CreateFAdd(srcl, dstl); - StoreInst *dsts = B.CreateStore(res, dsti); - StoreInst *srcs = nullptr; - if (zeroSrc) - srcs = B.CreateStore(Constant::getNullValue(res->getType()), srci); - if (dstalign) { - dsts->setAlignment(Align(dstalign)); - dstl->setAlignment(Align(dstalign)); - } - if (srcalign) { - if (zeroSrc) - srcs->setAlignment(Align(srcalign)); - srcl->setAlignment(Align(srcalign)); - } - - Value *nexti = - B.CreateAdd(i, ConstantInt::get(IT, 1), "i.next", true, true); - i->addIncoming(nexti, body); - B.CreateCondBr(B.CreateICmpEQ(nexti, M), initend, body); - } - - { - IRBuilder<> B(initend); - Value *nextj = - B.CreateAdd(j, ConstantInt::get(IT, 1), "j.next", true, true); - j->addIncoming(nextj, initend); - B.CreateCondBr(B.CreateICmpEQ(nextj, N), end, init); - } - - { - IRBuilder<> B(end); - B.CreateRetVoid(); - } - - return F; -} - // TODO implement differential memmove Function * getOrInsertDifferentialFloatMemmove(Module &M, Type *T, unsigned dstalign, diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index 83c9eccf5776..f8fff0fe4ad2 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -808,13 +808,8 @@ llvm::Function *getOrInsertMemcpyMat(llvm::Module &M, llvm::Type *elementType, llvm::Function *getOrInsertDifferentialFloatMemcpyMat( llvm::Module &M, llvm::Type *elementType, llvm::PointerType *PT, - llvm::IntegerType *IT, llvm::IntegerType *CT, unsigned dstalign, - unsigned srcalign, bool zeroSrc); - -llvm::Function *getOrInsertDifferentialFloatMemcpyMatLayout( - llvm::Module &M, llvm::Type *elementType, llvm::PointerType *PT, - llvm::IntegerType *IT, llvm::IntegerType *CT, unsigned dstalign, - unsigned srcalign, bool zeroSrc); + llvm::IntegerType *IT, llvm::IntegerType *UT, llvm::IntegerType *LT, + unsigned dstalign, unsigned srcalign, bool zeroSrc); /// Create function for type that performs the derivative memmove on floating /// point memory diff --git a/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll b/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll index fd2ce3a3ef0c..246d39ddbc08 100644 --- a/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll +++ b/enzyme/test/Enzyme/ForwardMode/blas/cblas_trsm_f.ll @@ -22,7 +22,7 @@ entry: ; CHECK-LABEL: define internal void @fwddiffef( ; CHECK: call void @cblas_dtrsm ; CHECK: icmp eq i32 %layout, 101 -; CHECK: __enzyme_dmemcpy_double_mat_layout_32_32_zero +; CHECK: __enzyme_dmemcpy_double_mat_8_32_32_zero ; CHECK: call void @cblas_dtrsm ; CHECK-NOT: call void bitcast -; CHECK: define internal void @__enzyme_dmemcpy_double_mat_layout_32_32_zero(i32 %layout, +; CHECK: define internal void @__enzyme_dmemcpy_double_mat_8_32_32_zero(i8 %uplo, i32 %layout, diff --git a/enzyme/test/Enzyme/ReverseMode/blas/dlacpy.ll b/enzyme/test/Enzyme/ReverseMode/blas/dlacpy.ll index 7c89fa56e598..3fa024056603 100644 --- a/enzyme/test/Enzyme/ReverseMode/blas/dlacpy.ll +++ b/enzyme/test/Enzyme/ReverseMode/blas/dlacpy.ll @@ -36,10 +36,7 @@ entry: ; CHECK-NEXT: %4 = load i64, i64* %ldb, align 8 ; CHECK-NEXT: call void @llvm.experimental.noalias.scope.decl(metadata !0) ; CHECK-NEXT: call void @llvm.experimental.noalias.scope.decl(metadata !3) -; CHECK-NEXT: %5 = icmp eq i64 %1, 0 -; CHECK-NEXT: %6 = icmp eq i64 %2, 0 -; CHECK-NEXT: %7 = or i1 %5, %6 -; CHECK-NEXT: br i1 %7, label %__enzyme_dmemcpy_double_mat_64_zero.exit, label %swtch.i +; CHECK: br i1 {{.*}}, label %__enzyme_dmemcpy_double_mat_8_8_64_zero.exit, label %swtch.i ; CHECK: swtch.i: ; preds = %invertentry ; CHECK-NEXT: switch i8 %0, label %Ginit.idx.i [ @@ -47,88 +44,19 @@ entry: ; CHECK-NEXT: i8 76, label %Linit.idx.i ; CHECK-NEXT: ] -; CHECK: Ginit.idx.i: ; preds = %Ginit.end.i, %swtch.i -; CHECK-NEXT: %Gj.i = phi i64 [ 0, %swtch.i ], [ %Gj.next.i, %Ginit.end.i ] -; CHECK-NEXT: br label %Gfor.body.i - -; CHECK: Uinit.idx.i: ; preds = %Uinit.end.i, %swtch.i -; CHECK-NEXT: %Uj.i = phi i64 [ 0, %swtch.i ], [ %Uj.next.i, %Uinit.end.i ] -; CHECK-NEXT: %8 = add nuw nsw i64 %Uj.i, 1 -; CHECK-NEXT: %9 = icmp ult i64 %8, %1 -; CHECK-NEXT: %10 = select i1 %9, i64 %8, i64 %1 -; CHECK-NEXT: br label %Ufor.body.i - -; CHECK: Linit.idx.i: ; preds = %Linit.end.i, %swtch.i -; CHECK-NEXT: %Lj.i = phi i64 [ 0, %swtch.i ], [ %Lj.next.i, %Linit.end.i ] -; CHECK-NEXT: br label %Lfor.body.i - ; CHECK: Gfor.body.i: ; preds = %Gfor.body.i, %Ginit.idx.i -; CHECK-NEXT: %Gi.i = phi i64 [ 0, %Ginit.idx.i ], [ %Gi.next.i, %Gfor.body.i ] -; CHECK-NEXT: %11 = mul nuw nsw i64 %Gj.i, %4 -; CHECK-NEXT: %12 = add nuw nsw i64 %Gi.i, %11 -; CHECK-NEXT: %Gsrc.i.i = getelementptr inbounds double, double* %"b'", i64 %12 -; CHECK-NEXT: %13 = mul nuw nsw i64 %Gj.i, %3 -; CHECK-NEXT: %14 = add nuw nsw i64 %Gi.i, %13 -; CHECK-NEXT: %Gdst.i.i = getelementptr inbounds double, double* %"a'", i64 %14 -; CHECK-NEXT: %Gsrc.i.l.i = load double, double* %Gsrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Gdst.i.l.i = load double, double* %Gdst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: %15 = fadd double %Gsrc.i.l.i, %Gdst.i.l.i -; CHECK-NEXT: store double %15, double* %Gdst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: store double 0.000000e+00, double* %Gsrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Gi.next.i = add nuw nsw i64 %Gi.i, 1 -; CHECK-NEXT: %16 = icmp eq i64 %Gi.next.i, %1 -; CHECK-NEXT: br i1 %16, label %Ginit.end.i, label %Gfor.body.i - -; CHECK: Ginit.end.i: ; preds = %Gfor.body.i -; CHECK-NEXT: %Gj.next.i = add nuw nsw i64 %Gj.i, 1 -; CHECK-NEXT: %17 = icmp eq i64 %Gj.next.i, %2 -; CHECK-NEXT: br i1 %17, label %__enzyme_dmemcpy_double_mat_64_zero.exit, label %Ginit.idx.i - +; CHECK: load double +; CHECK: fadd double +; CHECK: store double ; CHECK: Ufor.body.i: ; preds = %Ufor.body.i, %Uinit.idx.i -; CHECK-NEXT: %Ui.i = phi i64 [ 0, %Uinit.idx.i ], [ %Ui.next.i, %Ufor.body.i ] -; CHECK-NEXT: %18 = mul nuw nsw i64 %Uj.i, %4 -; CHECK-NEXT: %19 = add nuw nsw i64 %Ui.i, %18 -; CHECK-NEXT: %Usrc.i.i = getelementptr inbounds double, double* %"b'", i64 %19 -; CHECK-NEXT: %20 = mul nuw nsw i64 %Uj.i, %3 -; CHECK-NEXT: %21 = add nuw nsw i64 %Ui.i, %20 -; CHECK-NEXT: %Udst.i.i = getelementptr inbounds double, double* %"a'", i64 %21 -; CHECK-NEXT: %Usrc.i.l.i = load double, double* %Usrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Udst.i.l.i = load double, double* %Udst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: %22 = fadd double %Usrc.i.l.i, %Udst.i.l.i -; CHECK-NEXT: store double %22, double* %Udst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: store double 0.000000e+00, double* %Usrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Ui.next.i = add nuw nsw i64 %Ui.i, 1 -; CHECK-NEXT: %23 = icmp eq i64 %Ui.next.i, %10 -; CHECK-NEXT: br i1 %23, label %Uinit.end.i, label %Ufor.body.i - -; CHECK: Uinit.end.i: ; preds = %Ufor.body.i -; CHECK-NEXT: %Uj.next.i = add nuw nsw i64 %Uj.i, 1 -; CHECK-NEXT: %24 = icmp eq i64 %Uj.next.i, %2 -; CHECK-NEXT: br i1 %24, label %__enzyme_dmemcpy_double_mat_64_zero.exit, label %Uinit.idx.i - +; CHECK: load double +; CHECK: fadd double +; CHECK: store double ; CHECK: Lfor.body.i: ; preds = %Lfor.body.i, %Linit.idx.i -; CHECK-NEXT: %Li.i = phi i64 [ %Lj.i, %Linit.idx.i ], [ %Li.next.i, %Lfor.body.i ] -; CHECK-NEXT: %25 = mul nuw nsw i64 %Lj.i, %4 -; CHECK-NEXT: %26 = add nuw nsw i64 %Li.i, %25 -; CHECK-NEXT: %Lsrc.i.i = getelementptr inbounds double, double* %"b'", i64 %26 -; CHECK-NEXT: %27 = mul nuw nsw i64 %Lj.i, %3 -; CHECK-NEXT: %28 = add nuw nsw i64 %Li.i, %27 -; CHECK-NEXT: %Ldst.i.i = getelementptr inbounds double, double* %"a'", i64 %28 -; CHECK-NEXT: %Lsrc.i.l.i = load double, double* %Lsrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Ldst.i.l.i = load double, double* %Ldst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: %29 = fadd double %Lsrc.i.l.i, %Ldst.i.l.i -; CHECK-NEXT: store double %29, double* %Ldst.i.i, align 8, !alias.scope !0, !noalias !3 -; CHECK-NEXT: store double 0.000000e+00, double* %Lsrc.i.i, align 8, !alias.scope !3, !noalias !0 -; CHECK-NEXT: %Li.next.i = add nuw nsw i64 %Li.i, 1 -; CHECK-NEXT: %30 = icmp eq i64 %Li.next.i, %1 -; CHECK-NEXT: br i1 %30, label %Linit.end.i, label %Lfor.body.i - -; CHECK: Linit.end.i: ; preds = %Lfor.body.i -; CHECK-NEXT: %Lj.next.i = add nuw nsw i64 %Lj.i, 1 -; CHECK-NEXT: %31 = icmp eq i64 %Lj.next.i, %2 -; CHECK-NEXT: br i1 %31, label %__enzyme_dmemcpy_double_mat_64_zero.exit, label %Linit.idx.i - -; CHECK: __enzyme_dmemcpy_double_mat_64_zero.exit: ; preds = %invertentry, %Ginit.end.i, %Uinit.end.i, %Linit.end.i +; CHECK: load double +; CHECK: fadd double +; CHECK: store double +; CHECK: __enzyme_dmemcpy_double_mat_8_8_64_zero.exit: ; preds = ; CHECK-NEXT: ret void ; CHECK-NEXT: } @@ -178,4 +106,3 @@ entry: ; CHECK-NEXT: call void @dlascl_64_(i8* %uplo, i64* %byref.constant.int.0, i64* %byref.constant.int.01, double* %byref.constant.fp.1.0, double* %byref.constant.fp.0.0, i64* %m, i64* %n, double* %"b'", i64* %ldb, i64* %0, i64 1) ; CHECK-NEXT: ret void ; CHECK-NEXT: } - diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp index eb68eb89c950..9c40a27d9cd2 100644 --- a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp @@ -2004,20 +2004,34 @@ void emit_dag(bool forward, Twine resultVarName, const DagInit *ruleDag, << ValueType_helper(pattern, actArg, ruleDag) << "}, Builder2, /* lookup */ " << (!forward) << ");\n"; - for (int i : {0, 1, 2, 4, 6}) - os << " " << argPrefix << "[" << i << "] = load_if_ref(Builder2, " - << ((i == 0) ? "charType" : "intType") << ", " << argPrefix << "[" << i - << "], byRef);\n"; + os << " Value *layout = cblas ? load_if_ref(Builder2, charType, " + "arg_layout, byRef) : ConstantInt::get(charType, 102);\n"; + os << " unsigned uploOffset = cblas ? 1 : 0;\n"; + os << " auto uplo = load_if_ref(Builder2, charType, " << argPrefix + << "[uploOffset], byRef);\n"; + os << " auto M = load_if_ref(Builder2, intType, " << argPrefix + << "[uploOffset + 1], byRef);\n"; + os << " auto N = load_if_ref(Builder2, intType, " << argPrefix + << "[uploOffset + 2], byRef);\n"; + os << " auto dstLD = load_if_ref(Builder2, intType, " << argPrefix + << "[uploOffset + 4], byRef);\n"; + os << " auto srcLD = load_if_ref(Builder2, intType, " << argPrefix + << "[uploOffset + 6], byRef);\n"; os << " auto dmemcpymat = " "getOrInsertDifferentialFloatMemcpyMat(*gutils->oldFunc->getParent()," " fpType, cast(" - << argPrefix << "[" << argPrefix - << ".size() - 2]->getType()), intType, charType, 0, 0, " + << argPrefix + << "[uploOffset + 3]->getType()), intType, " + "cast(uplo->getType()), " + "cast(layout->getType()), 0, 0, " << Def->getValueAsBit("shouldZero") << ");\n"; + os << " Value *matadd_args[] = {uplo, layout, M, N, " << argPrefix + << "[uploOffset + 3], dstLD, " << argPrefix + << "[uploOffset + 5], srcLD};\n"; os << " auto cubcall = cast(Builder2.CreateCall(dmemcpymat, " - << argPrefix << ", Defs));\n"; + "matadd_args, Defs));\n"; if (!forward && !runtimeChecked) emit_runtime_continue(ruleDag, argName, " ", "Builder2", @@ -2059,12 +2073,15 @@ void emit_dag(bool forward, Twine resultVarName, const DagInit *ruleDag, << "_3[1], byRef);\n"; os << " auto srcLD = load_if_ref(Builder2, intType, " << argPrefix << "_4[1], byRef);\n"; + os << " Value *uplo = ConstantInt::get(charType, 'G');\n"; + os << " auto uploType = cast(uplo->getType());\n"; os << " auto dmemcpymat = " - "getOrInsertDifferentialFloatMemcpyMatLayout(*gutils->oldFunc->" + "getOrInsertDifferentialFloatMemcpyMat(*gutils->oldFunc->" "getParent(), fpType, cast(" - << argPrefix << "_3[0]->getType()), intType, layoutType, 0, 0, " + << argPrefix + << "_3[0]->getType()), intType, uploType, layoutType, 0, 0, " << Def->getValueAsBit("shouldZero") << ");\n"; - os << " Value *matadd_args[] = {layout, M, N, " << argPrefix + os << " Value *matadd_args[] = {uplo, layout, M, N, " << argPrefix << "_3[0], dstLD, " << argPrefix << "_4[0], srcLD};\n"; os << " auto cubcall = cast(Builder2.CreateCall(" "dmemcpymat, matadd_args, Defs));\n";