Add complex sqrt (#1778)

* Add complex sqrt * fix analyzer capture * fix analyzer error passing convention * Handle extractions in tablegen
EnzymeAD · Apr 1, 2024 · 5b2f04c · 5b2f04c
1 parent b1f676c
commit 5b2f04c
Show file tree

Hide file tree

Showing 9 changed files with 174 additions and 42 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -455,7 +455,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
       }
       if (CustomErrorHandler) {
         CustomErrorHandler(str.c_str(), wrap(&I), ErrorType::NoType,
-                           &TR.analyzer, nullptr, wrap(&BuilderZ));
+                           TR.analyzer, nullptr, wrap(&BuilderZ));
       } else {
         ss << "\n";
         TR.dump(ss);
@@ -1032,7 +1032,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
       }
       if (CustomErrorHandler) {
         CustomErrorHandler(str.c_str(), wrap(&I), ErrorType::NoType,
-                           &TR.analyzer, nullptr, wrap(&BuilderZ));
+                           TR.analyzer, nullptr, wrap(&BuilderZ));
       } else {
         ss << "\n";
         TR.dump(ss);
@@ -1058,7 +1058,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                << " size: " << storeSize;
             if (CustomErrorHandler) {
               CustomErrorHandler(str.c_str(), wrap(&I), ErrorType::NoType,
-                                 &TR.analyzer, nullptr, wrap(&BuilderZ));
+                                 TR.analyzer, nullptr, wrap(&BuilderZ));
             } else {
               EmitFailure("CannotDeduceType", I.getDebugLoc(), &I, ss.str());
             }
@@ -1136,7 +1136,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
            << " size: " << storeSize;
         if (CustomErrorHandler) {
           CustomErrorHandler(str.c_str(), wrap(&I), ErrorType::NoType,
-                             &TR.analyzer, nullptr, wrap(&BuilderZ));
+                             TR.analyzer, nullptr, wrap(&BuilderZ));
         } else {
           EmitFailure("CannotDeduceType", I.getDebugLoc(), &I, ss.str());
         }
@@ -1445,7 +1445,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
           ss << "Cannot deduce adding type (cast) of " << I;
           if (CustomErrorHandler) {
             CustomErrorHandler(ss.str().c_str(), wrap(&I), ErrorType::NoType,
-                               &TR.analyzer, nullptr, wrap(&Builder2));
+                               TR.analyzer, nullptr, wrap(&Builder2));
             return;
           } else {
             ss << "\n";
@@ -2012,7 +2012,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                  << " size: " << size0 << " TT: " << TT.str();
               if (CustomErrorHandler) {
                 CustomErrorHandler(str.c_str(), wrap(&IVI), ErrorType::NoType,
-                                   &TR.analyzer, nullptr, wrap(&Builder2));
+                                   TR.analyzer, nullptr, wrap(&Builder2));
               } else {
                 EmitFailure("CannotDeduceType", IVI.getDebugLoc(), &IVI,
                             ss.str());
@@ -2097,7 +2097,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                  << " TT: " << TT.str();
               if (CustomErrorHandler) {
                 CustomErrorHandler(str.c_str(), wrap(&IVI), ErrorType::NoType,
-                                   &TR.analyzer, nullptr, wrap(&Builder2));
+                                   TR.analyzer, nullptr, wrap(&Builder2));
               } else {
                 EmitFailure("CannotDeduceType", IVI.getDebugLoc(), &IVI,
                             ss.str());
@@ -3133,7 +3133,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
       ss << "Cannot deduce type of memset " << MS;
       if (CustomErrorHandler) {
         CustomErrorHandler(str.c_str(), wrap(&MS), ErrorType::NoType,
-                           &TR.analyzer, nullptr, wrap(&BuilderZ));
+                           TR.analyzer, nullptr, wrap(&BuilderZ));
       } else {
         ss << "\n";
         TR.dump(ss);
@@ -3439,7 +3439,7 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
         ss << "Cannot deduce type of copy " << MTI;
         if (CustomErrorHandler) {
           CustomErrorHandler(str.c_str(), wrap(&MTI), ErrorType::NoType,
-                             &TR.analyzer, nullptr, wrap(&BuilderZ));
+                             TR.analyzer, nullptr, wrap(&BuilderZ));
         } else {
           ss << "\n";
           ss << *gutils->oldFunc << "\n";

diff --git a/enzyme/Enzyme/DiffeGradientUtils.cpp b/enzyme/Enzyme/DiffeGradientUtils.cpp
@@ -496,7 +496,7 @@ DiffeGradientUtils::addToDiffe(Value *val, Value *dif, IRBuilder<> &BuilderM,
       ss << "}\n";
       if (CustomErrorHandler) {
         CustomErrorHandler(ss.str().c_str(), wrap(val), ErrorType::NoType,
-                           &TR.analyzer, nullptr, wrap(&BuilderM));
+                           TR.analyzer, nullptr, wrap(&BuilderM));
         return addedSelects;
       } else {
         TR.dump(ss);
@@ -534,7 +534,7 @@ DiffeGradientUtils::addToDiffe(Value *val, Value *dif, IRBuilder<> &BuilderM,
          << "\n";
       if (CustomErrorHandler) {
         CustomErrorHandler(ss.str().c_str(), wrap(val), ErrorType::NoType,
-                           &TR.analyzer, nullptr, wrap(&BuilderM));
+                           TR.analyzer, nullptr, wrap(&BuilderM));
         return addedSelects;
       } else {
         DebugLoc loc;

diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
@@ -1802,6 +1802,16 @@ class EnzymeBase {
           csts.push_back(ConstantFP::get(e, 1.0));
         }
         args.push_back(ConstantStruct::get(ST, csts));
+      } else if (auto AT = dyn_cast<ArrayType>(fn->getReturnType())) {
+        SmallVector<Constant *, 2> csts(
+            AT->getNumElements(), ConstantFP::get(AT->getElementType(), 1.0));
+        args.push_back(ConstantArray::get(AT, csts));
+      } else {
+        auto RT = fn->getReturnType();
+        EmitFailure("EnzymeCallingError", CI->getDebugLoc(), CI,
+                    "Differential return required for call ", *CI,
+                    " but one of type ", *RT, " could not be auto deduced");
+        return false;
       }
     }
 

diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -3364,7 +3364,7 @@ void createInvertedTerminator(DiffeGradientUtils *gutils,
          << " sz: " << size << "\n";
       if (CustomErrorHandler) {
         CustomErrorHandler(str.c_str(), wrap(orig), ErrorType::NoType,
-                           &gutils->TR.analyzer, nullptr, wrap(&Builder));
+                           gutils->TR.analyzer, nullptr, wrap(&Builder));
         continue;
       } else {
         ss << "\n";

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -4907,6 +4907,21 @@ Type *GradientUtils::getShadowType(Type *ty) {
   return getShadowType(ty, width);
 }
 
+Type *GradientUtils::extractMeta(Type *T, ArrayRef<unsigned> off) {
+  for (auto idx : off) {
+    if (auto AT = dyn_cast<ArrayType>(T)) {
+      T = AT->getElementType();
+      continue;
+    }
+    if (auto ST = dyn_cast<StructType>(T)) {
+      T = ST->getElementType(idx);
+      continue;
+    }
+    assert(false && "could not sub index into type");
+  }
+  return T;
+}
+
 Value *GradientUtils::extractMeta(IRBuilder<> &Builder, Value *Agg,
                                   unsigned off, const Twine &name) {
   return extractMeta(Builder, Agg, ArrayRef<unsigned>({off}), name);

diff --git a/enzyme/Enzyme/GradientUtils.h b/enzyme/Enzyme/GradientUtils.h
@@ -538,6 +538,9 @@ class GradientUtils : public CacheUtility {
                                   const llvm::Twine &name = "",
                                   bool fallback = true);
 
+  //! Helper routine to get the type of an extraction
+  static llvm::Type *extractMeta(llvm::Type *T, llvm::ArrayRef<unsigned> off);
+
   static llvm::Value *recursiveFAdd(llvm::IRBuilder<> &B, llvm::Value *lhs,
                                     llvm::Value *rhs,
                                     llvm::ArrayRef<unsigned> lhs_off = {},

diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -934,6 +934,13 @@ def : IntrPattern<(Op $x),
                   (ForwardFromSummedReverse)
                  >;
 
+
+def : CallPattern<(Op (Op $x, $y):$z),
+                  ["cmplx_sqrt"],
+                  [(Select (And (FCmpUEQ $x, (ConstantFP<"0"> $x)), (FCmpUEQ $y, (ConstantFP<"0"> $y))), (ConstantCFP<"0", "0"> $z), (Conj (CFDiv (Conj (DiffeRet)), (CFMul (ConstantCFP<"2", "0"> $z), (Call<(SameFunc), [ReadNone,NoUnwind]> $z)))))],
+                  (Select (And (FCmpUEQ $x, (ConstantFP<"0"> $x)), (FCmpUEQ $y, (ConstantFP<"0"> $y))), (ConstantCFP<"0", "0"> $z), (CFDiv (Shadow $z), (CFMul (ConstantCFP<"2", "0"> $z), (Call<(SameFunc), [ReadNone,NoUnwind]> $z))))
+                 >;
+
 def : IntrPattern<(Op $x, $y),
                   [["pow"]],
                   [

diff --git a/enzyme/test/Enzyme/ReverseMode/cmplx_sqrt.ll b/enzyme/test/Enzyme/ReverseMode/cmplx_sqrt.ll
@@ -0,0 +1,20 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s
+
+declare [2 x double] @cmplx_sqrt([2 x double] %x)
+
+define [2 x double] @tester([2 x double] %x) {
+entry:
+  %y = call [2 x double] @cmplx_sqrt([2 x double] %x)
+  ret [2 x double] %y
+}
+
+define [2 x double] @test_derivative([2 x double] %x) {
+entry:
+  %0 = tail call [2 x double] (...) @__enzyme_autodiff([2 x double] ([2 x double])* nonnull @tester, metadata !"enzyme_active_return", [2 x double] %x)
+  ret [2 x double] %0
+}
+
+declare [2 x double] @__enzyme_autodiff(...)
+
+; CHECK: define internal { [2 x double] } @diffetester([2 x double] %x, [2 x double] %differeturn)