Fix zero-sized cholesky (#2157)

* Fix zero-sized cholesky * fix byref
EnzymeAD · Nov 6, 2024 · 342057e · 342057e
1 parent f9b6768
commit 342057e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/enzyme/Enzyme/BlasDerivatives.td b/enzyme/Enzyme/BlasDerivatives.td
@@ -101,6 +101,7 @@ def is_diag_int : IntMatchers<
 
 def is_left : MagicInst;
 def is_lower : MagicInst;
+def is_zero : MagicInst;
 def is_nonunit : MagicInst;
 
 def First : MagicInst;
@@ -757,7 +758,7 @@ def potrf: CallBlasPattern<(Op $layout, $uplo, $n, $A, $lda, $info),
 
                       (BlasCall<"trsm"> $layout, (uplo_to_side $uplo), $uplo, Char<"T">, Char<"N">, $n, $n, Constant<"1.0">, $A, (ld $A, Char<"N">, $lda, $n, $n), use<"ztri">, $n) ,
 
-                      (For<"i", 0> (Sub $n, ConstantInt<1>),
+                      (For<"i", 0> (ISelect (is_zero $n), ConstantInt<0>, (Sub $n, ConstantInt<1>)),
                         (BlasCall<"axpy">
                             (Sub (Sub $n, ConstantInt<1>), $i),
                             Constant<"1.0">,

diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp
@@ -928,6 +928,20 @@ void rev_call_arg(bool forward, const DagInit *ruleDag,
            << ", cache_" << matName << ", byRef, cublas)}";
         return;
       }
+      if (Def->getName() == "is_zero") {
+        if (Dag->getNumArgs() != 1)
+          PrintFatalError(pattern.getLoc(), "only 1-arg ld operands supported");
+        const auto name = Dag->getArgNameStr(0);
+        os << "    ({ auto V = load_if_ref(Builder2, intType, arg_" << name
+           << ", byRef);\n";
+        os << "    SmallVector<Value*, 1> vs = {to_blas_callconv(Builder2, "
+              "Builder2.CreateICmpEQ(V, ConstantInt::get(V->getType(), 0)), "
+              "byRef, cublas, julia_decl_type, allocationBuilder, "
+              "\"is_zero\")};\n";
+        os << "    vs; })";
+        return;
+      }
+
       if (Def->getName() == "is_left") {
         if (Dag->getNumArgs() != 1)
           PrintFatalError(pattern.getLoc(), "only 1-arg ld operands supported");