From 2105ebc6dd0ab09ee785633ca936dc8f35ee335d Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Sat, 16 May 2026 15:55:50 -0700 Subject: [PATCH] feat: add BLAS trsv derivative rule Adds the `def trsv` rule for the real triangular vector solve, with both forward and reverse mode derivatives. Mirrors the structure of `trtrs`. Outside the new `def`, the only required tablegen additions are: - `def Dep : MagicInst;` with a 6-line passthrough in `rev_call_arg`. The forward rule needs to gate two `axpy` and one `copy` call on `Shadow $A` being active, but `Shadow $A` is not a direct BLAS arg of those calls. `Dep` wraps an arg so activity analysis sees the dependency without it being emitted as a real BLAS argument. - `trsv` joins `trmv`/`trtrs` in the n=3 byRef-suffix arity table (three trailing string-length args for uplo/trans/diag). - `trsv` joins `trtrs` in the reversed activeArg iteration order, since reverse mode needs the solved x cotangent before forming dA. No changes to Utils.cpp/Utils.h or to other BLAS rules. --- enzyme/Enzyme/BlasDerivatives.td | 40 ++++++++++++++++++ enzyme/test/Enzyme/ForwardMode/blas/trsv_f.ll | 42 +++++++++++++++++++ enzyme/test/Enzyme/ReverseMode/blas/trsv_f.ll | 41 ++++++++++++++++++ enzyme/tools/enzyme-tblgen/blas-tblgen.cpp | 15 +++++-- 4 files changed, 134 insertions(+), 4 deletions(-) create mode 100644 enzyme/test/Enzyme/ForwardMode/blas/trsv_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..411e3c1d3d3c 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; @@ -672,6 +673,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), $x), use<"tmp">, ConstantInt<1>), + (BlasCall<"trmv"> $layout, $uplo, $trans, $diag, $n, (Shadow $A), 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"]>], 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/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/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp index dbdcc0f6a6f6..0fe850612677 100644 --- a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp @@ -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"); @@ -1493,7 +1499,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; @@ -2354,9 +2360,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); + // trtrs and trsv solve x 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];