Generate code for Enzyme autodiff for functions with Pointer/array ar…

…guments For functions of type double myfunction(double* arr){...}; OR double myfunction(double arr[3]){...}; we generate derivative code of type: void d_myfunction(double* arr, clad::array_ref<double> _d_arr){ double* d_arr = _d_arr.ptr(); __enzyme_autodiff_myfunction(myfunction, arr, d_arr); } Further on enzyme will handle the function differentiation This Commit also Downloads and Installs Enzyme in the Clad Directory, if the flag "-DENABLE_ENZYME_BACKEND=On" is passed to cmake during project configuration. This commit also contains tests for the above feature in test/Gradient/Enzyme.C This commit also adds the enzyme pass to the llvm optimisation pipeline so that it can run at the right moment.
vgvassilev · Aug 15, 2022 · 828baf7 · 828baf7
1 parent b5bfdb9
commit 828baf7
Show file tree

Hide file tree

Showing 13 changed files with 283 additions and 85 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -62,13 +62,13 @@ jobs:
             compiler: gcc-5
             clang-runtime: '8'
 
-          - name: ubu18-gcc6-runtime11-analyzers
+          - name: ubu18-gcc7-runtime11-analyzers
             os: ubuntu-18.04
-            compiler: gcc-6
+            compiler: gcc-7
             clang-runtime: '11'
             coverage: true
             cuda: true
-            extra_cmake_options: '-DLLVM_ENABLE_WERROR=On'
+            extra_cmake_options: '-DLLVM_ENABLE_WERROR=On -DENABLE_ENZYME_BACKEND=On'
             #clang-format: true
 
           - name: ubu18-gcc6-runtime11-benchmarks

diff --git a/.gitignore b/.gitignore
@@ -1,2 +1,2 @@
 /build
-/.vscode
+/.vscode
diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h
@@ -86,6 +86,11 @@ namespace clad {
         quals.removeConst();
         return S.BuildQualifiedType(T.getUnqualifiedType(), noLoc, quals);
     }
+    // Function to Differentiate with Clad as Backend
+    void DifferentiateWithClad();
+
+    // Function to Differentiate with Enzyme as Backend
+    void DifferentiateWithEnzyme();
 
   public:
     using direction = rmv::direction;

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -408,69 +408,12 @@ namespace clad {
 
     Stmt* gradientBody = nullptr;
 
-    // create derived variables for parameters which are not part of
-    // independent variables (args).
-    if (!use_enzyme) {
-      for (std::size_t i = 0; i < m_Function->getNumParams(); ++i) {
-        ParmVarDecl* param = params[i];
-        // derived variables are already created for independent variables.
-        if (m_Variables.count(param))
-          continue;
-        // in vector mode last non diff parameter is output parameter.
-        if (isVectorValued && i == m_Function->getNumParams() - 1)
-          continue;
-        auto VDDerivedType = param->getType();
-        // We cannot initialize derived variable for pointer types because
-        // we do not know the correct size.
-        if (utils::isArrayOrPointerType(VDDerivedType))
-          continue;
-        auto VDDerived =
-            BuildVarDecl(VDDerivedType, "_d_" + param->getNameAsString(),
-                         getZeroInit(VDDerivedType));
-        m_Variables[param] = BuildDeclRef(VDDerived);
-        addToBlock(BuildDeclStmt(VDDerived), m_Globals);
-      }
-      // Start the visitation process which outputs the statements in the
-      // current block.
-      StmtDiff BodyDiff = Visit(FD->getBody());
-      Stmt* Forward = BodyDiff.getStmt();
-      Stmt* Reverse = BodyDiff.getStmt_dx();
-      // Create the body of the function.
-      // Firstly, all "global" Stmts are put into fn's body.
-      for (Stmt* S : m_Globals)
-        addToCurrentBlock(S, direction::forward);
-      // Forward pass.
-      if (auto CS = dyn_cast<CompoundStmt>(Forward))
-        for (Stmt* S : CS->body())
-          addToCurrentBlock(S, direction::forward);
-      else
-        addToCurrentBlock(Forward, direction::forward);
-      // Reverse pass.
-      if (auto RCS = dyn_cast<CompoundStmt>(Reverse))
-        for (Stmt* S : RCS->body())
-          addToCurrentBlock(S, direction::forward);
-      else
-        addToCurrentBlock(Reverse, direction::forward);
-
-      if (m_ExternalSource)
-        m_ExternalSource->ActOnEndOfDerivedFnBody();
-
-      gradientBody = endBlock();
-    } else {
-      // TODO: Write code to generate code that enzyme can work on
-      /*
-        Two if cases ought to come here
-        a. How to deal with functions of form - double function(double z, double
-        y, double z...); This is not straight forward as of now because Enzyme
-        does not have support for this yet, we will need to setup data
-        structures that can deal with this easily
-
-        b. How to deal with functions of the form - double function(double*
-        arrayOfParams); This is straightforward to deal in enzyme
-      */
-      gradientBody = endBlock();
-    }
+    if (!use_enzyme)
+      DifferentiateWithClad();
+    else
+      DifferentiateWithEnzyme();
 
+    gradientBody = endBlock();
     m_Derivative->setBody(gradientBody);
     endScope(); // Function body scope
     m_Sema.PopFunctionScopeInfo();
@@ -590,6 +533,121 @@ namespace clad {
     return DerivativeAndOverload{fnBuildRes.first, nullptr};
   }
 
+  void ReverseModeVisitor::DifferentiateWithClad() {
+    llvm::ArrayRef<ParmVarDecl*> paramsRef = m_Derivative->parameters();
+
+    // create derived variables for parameters which are not part of
+    // independent variables (args).
+    for (std::size_t i = 0; i < m_Function->getNumParams(); ++i) {
+      ParmVarDecl* param = paramsRef[i];
+      // derived variables are already created for independent variables.
+      if (m_Variables.count(param))
+        continue;
+      // in vector mode last non diff parameter is output parameter.
+      if (isVectorValued && i == m_Function->getNumParams() - 1)
+        continue;
+      auto VDDerivedType = param->getType();
+      // We cannot initialize derived variable for pointer types because
+      // we do not know the correct size.
+      if (utils::isArrayOrPointerType(VDDerivedType))
+        continue;
+      auto VDDerived =
+          BuildVarDecl(VDDerivedType, "_d_" + param->getNameAsString(),
+                       getZeroInit(VDDerivedType));
+      m_Variables[param] = BuildDeclRef(VDDerived);
+      addToBlock(BuildDeclStmt(VDDerived), m_Globals);
+    }
+    // Start the visitation process which outputs the statements in the
+    // current block.
+    StmtDiff BodyDiff = Visit(m_Function->getBody());
+    Stmt* Forward = BodyDiff.getStmt();
+    Stmt* Reverse = BodyDiff.getStmt_dx();
+    // Create the body of the function.
+    // Firstly, all "global" Stmts are put into fn's body.
+    for (Stmt* S : m_Globals)
+      addToCurrentBlock(S, direction::forward);
+    // Forward pass.
+    if (auto CS = dyn_cast<CompoundStmt>(Forward))
+      for (Stmt* S : CS->body())
+        addToCurrentBlock(S, direction::forward);
+    else
+      addToCurrentBlock(Forward, direction::forward);
+    // Reverse pass.
+    if (auto RCS = dyn_cast<CompoundStmt>(Reverse))
+      for (Stmt* S : RCS->body())
+        addToCurrentBlock(S, direction::forward);
+    else
+      addToCurrentBlock(Reverse, direction::forward);
+
+    if (m_ExternalSource)
+      m_ExternalSource->ActOnEndOfDerivedFnBody();
+  }
+
+  void ReverseModeVisitor::DifferentiateWithEnzyme() {
+    // FIXME: Generalize this function to differentiate other kinds
+    // of function prototypes
+    unsigned numParams = m_Function->getNumParams();
+    auto origParams = m_Function->parameters();
+    llvm::ArrayRef<ParmVarDecl*> paramsRef = m_Derivative->parameters();
+    auto originalFnType = dyn_cast<FunctionProtoType>(m_Function->getType());
+
+    // Case 1: The function to be differentiated is of type double
+    // func(double* arr){...};
+    //        or double func(double arr[n]){...};
+    if (numParams == 1) {
+      QualType origTy = origParams[0]->getOriginalType();
+      if (origTy->isConstantArrayType() || origTy->isPointerType()) {
+        // Extract Pointer from Clad Array Ref
+        auto arrayRefNameExpr = BuildDeclRef(paramsRef[1]);
+        auto getPointerExpr = BuildCallExprToMemFn(arrayRefNameExpr, "ptr", {});
+        auto arrayRefToArrayStmt = BuildVarDecl(
+            origTy, "d_" + paramsRef[0]->getNameAsString(), getPointerExpr);
+        addToCurrentBlock(BuildDeclStmt(arrayRefToArrayStmt),
+                          direction::forward);
+
+        // Prepare Arguments to enzyme_autodiff
+        llvm::SmallVector<Expr*, 16> enzymeArgs;
+        enzymeArgs.push_back(
+            BuildDeclRef(const_cast<FunctionDecl*>(m_Function)));
+        enzymeArgs.push_back(BuildDeclRef(paramsRef[0]));
+        enzymeArgs.push_back(BuildDeclRef(arrayRefToArrayStmt));
+
+        // Prepare Parameters for Function Signature
+        llvm::SmallVector<ParmVarDecl*, 16> enzymeParams;
+        DeclContext* fdDeclContext =
+            const_cast<DeclContext*>(m_Function->getDeclContext());
+        enzymeParams.push_back(m_Sema.BuildParmVarDeclForTypedef(
+            fdDeclContext, noLoc, m_Function->getType()));
+        enzymeParams.push_back(m_Sema.BuildParmVarDeclForTypedef(
+            fdDeclContext, noLoc, paramsRef[0]->getType()));
+        enzymeParams.push_back(m_Sema.BuildParmVarDeclForTypedef(
+            fdDeclContext, noLoc, arrayRefToArrayStmt->getType()));
+
+        // Get the Parameter Types in Function Signature
+        llvm::SmallVector<QualType, 8> enzymeParamsType;
+        for (auto i : enzymeParams)
+          enzymeParamsType.push_back(i->getType());
+
+        // Prepare Function call
+        std::string enzymeCallName =
+            "__enzyme_autodiff_" + m_Function->getNameAsString();
+        IdentifierInfo* IIEnzyme = &m_Context.Idents.get(enzymeCallName);
+        DeclarationName nameEnzyme(IIEnzyme);
+        QualType enzymeFunctionType = m_Sema.BuildFunctionType(
+            m_Context.VoidTy, enzymeParamsType, noLoc, nameEnzyme,
+            originalFnType->getExtProtoInfo());
+        FunctionDecl* enzymeCallFD = FunctionDecl::Create(
+            m_Context, const_cast<DeclContext*>(m_Function->getDeclContext()),
+            noLoc, noLoc, nameEnzyme, enzymeFunctionType,
+            m_Function->getTypeSourceInfo(), SC_Extern);
+        enzymeCallFD->setParams(enzymeParams);
+
+        // Add Function call to block
+        Expr* enzymeCall = BuildCallExprToFunction(enzymeCallFD, enzymeArgs);
+        addToCurrentBlock(enzymeCall);
+      }
+    }
+  }
   StmtDiff ReverseModeVisitor::VisitStmt(const Stmt* S) {
     diag(
         DiagnosticsEngine::Warning,

diff --git a/patches/enzyme.patch b/patches/enzyme.patch
@@ -0,0 +1,78 @@
+From b2451b5910d556c291330ca4d51903b57537b2e0 Mon Sep 17 00:00:00 2001
+From: Vassil Vassilev <[email protected]>
+Date: Fri, 29 Jul 2022 08:30:24 +0000
+Subject: [PATCH] Simplify integration in Clad
+
+---
+ enzyme/CMakeLists.txt        | 22 +++++++++++-----------
+ enzyme/Enzyme/CMakeLists.txt |  9 +++++++++
+ 2 files changed, 20 insertions(+), 11 deletions(-)
+
+diff --git a/enzyme/CMakeLists.txt b/enzyme/CMakeLists.txt
+index 4095147..09edfb0 100644
+--- a/enzyme/CMakeLists.txt
++++ b/enzyme/CMakeLists.txt
+@@ -181,7 +181,7 @@ add_subdirectory(test)
+
+ # The benchmarks data are not in git-exported source archives to minimize size.
+ # Only add the benchmarks if the directory exists.
+-if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/benchmarks")
++if (ENZYME_ENABLE_BENCHMARKS)
+     add_subdirectory(benchmarks)
+ endif()
+
+@@ -202,8 +202,8 @@ export(TARGETS ClangEnzyme-${LLVM_VERSION_MAJOR}
+     APPEND FILE "${PROJECT_BINARY_DIR}/EnzymeTargets.cmake")
+ endif()
+
+-export(TARGETS LLDEnzyme-${LLVM_VERSION_MAJOR}
+-    APPEND FILE "${PROJECT_BINARY_DIR}/EnzymeTargets.cmake")
++#export(TARGETS LLDEnzyme-${LLVM_VERSION_MAJOR}
++#    APPEND FILE "${PROJECT_BINARY_DIR}/EnzymeTargets.cmake")
+
+ export(PACKAGE Enzyme)
+
+@@ -223,12 +223,12 @@ configure_file(cmake/EnzymeConfig.cmake.in
+
+ configure_file(cmake/EnzymeConfigVersion.cmake.in
+      "${PROJECT_BINARY_DIR}/EnzymeConfigVersion.cmake" @ONLY)
+-   
+-install(FILES
+-     "${PROJECT_BINARY_DIR}${CMAKE_FILES_DIRECTORY}/EnzymeConfig.cmake"
+-     "${PROJECT_BINARY_DIR}/EnzymeConfigVersion.cmake"
+-     DESTINATION "${INSTALL_CMAKE_DIR}" COMPONENT dev)
+-   
+-install(EXPORT EnzymeTargets DESTINATION
+-    "${INSTALL_CMAKE_DIR}" COMPONENT dev)
++
++#install(FILES
++#     "${PROJECT_BINARY_DIR}${CMAKE_FILES_DIRECTORY}/EnzymeConfig.cmake"
++#     "${PROJECT_BINARY_DIR}/EnzymeConfigVersion.cmake"
++#     DESTINATION "${INSTALL_CMAKE_DIR}" COMPONENT dev)
++#
++#install(EXPORT EnzymeTargets DESTINATION
++#    "${INSTALL_CMAKE_DIR}" COMPONENT dev)
+
+diff --git a/enzyme/Enzyme/CMakeLists.txt b/enzyme/Enzyme/CMakeLists.txt
+index 78891b5..13fdb9c 100644
+--- a/enzyme/Enzyme/CMakeLists.txt
++++ b/enzyme/Enzyme/CMakeLists.txt
+@@ -13,6 +13,15 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
+ list(APPEND ENZYME_SRC SCEV/ScalarEvolutionExpander.cpp)
+ list(APPEND ENZYME_SRC  TypeAnalysis/TypeTree.cpp TypeAnalysis/TypeAnalysis.cpp TypeAnalysis/TypeAnalysisPrinter.cpp TypeAnalysis/RustDebugInfo.cpp)
+
++if (ENZYME_BUILD_STATIC_ONLY)
++  add_llvm_library( LLVMEnzyme-${LLVM_VERSION_MAJOR}
++    ${ENZYME_SRC}
++    DEPENDS
++    intrinsics_gen
++    )
++  return()
++endif()
++
+ if (${LLVM_VERSION_MAJOR} LESS 8)
+     add_llvm_loadable_module( LLVMEnzyme-${LLVM_VERSION_MAJOR}
+         ${ENZYME_SRC}
+-- 
+2.17.1
+
diff --git a/test/ForwardMode/Enzyme.C → test/Enzyme/ForwardMode.C b/test/ForwardMode/Enzyme.C → test/Enzyme/ForwardMode.C
@@ -1,13 +1,15 @@
 // RUN: %cladclang %s -lstdc++ -I%S/../../include -oEnzyme.out 2>&1 | FileCheck %s
 // RUN: ./Enzyme.out
 // CHECK-NOT: {{.*error|warning|note:.*}}
+// REQUIRES: Enzyme
 // XFAIL:*
+// Forward mode is not implemented yet
 
 #include "clad/Differentiator/Differentiator.h"
 
 double f(double x, double y) { return x * y; }
 
-// CHECK:     void f_diff_enzyme() {
+// CHECK:     double f_darg0_enzyme(double x, double y) {
 // CHECK-NEXT:}
 
 int main(){

diff --git a/test/Enzyme/ReverseMode.C b/test/Enzyme/ReverseMode.C
@@ -0,0 +1,21 @@
+// RUN: %cladclang %s -I%S/../../include -oReverseMode.out | FileCheck %s
+// RUN: ./ReverseMode.out | FileCheck -check-prefix=CHECK-EXEC %s
+// CHECK-NOT: {{.*error|warning|note:.*}}
+// REQUIRES: Enzyme
+
+#include "clad/Differentiator/Differentiator.h"
+
+double f(double* arr) { return arr[0] * arr[1]; }
+
+// CHECK: void f_grad_enzyme(double *arr, clad::array_ref<double> _d_arr) {
+// CHECK-NEXT:    double *d_arr = _d_arr.ptr();
+// CHECK-NEXT:    __enzyme_autodiff_f(f, arr, d_arr);
+// CHECK-NEXT:}
+
+int main() {
+  auto f_grad = clad::gradient<clad::opts::use_enzyme>(f);
+  double v[2] = {3, 4};
+  double g[2] = {0};
+  f_grad.execute(v, g);
+  printf("d_x = %.2f, d_y = %.2f", g[0], g[1]); // CHECK-EXEC: d_x = 4.00, d_y = 3.00
+}
diff --git a/test/Gradient/Enzyme.C b/test/Gradient/Enzyme.C
diff --git a/test/lit.cfg b/test/lit.cfg
@@ -245,11 +245,10 @@ def inferClang(PATH):
     return clang
 
 config.clang = inferClang(config.environment['PATH']).replace('\\', '/')
-
 config.cladlib = inferCladLib(config.environment['PATH']).replace('\\', '/')
+
 if not lit_config.quiet:
     lit_config.note('using clang: %r' % config.clang)
-if not lit_config.quiet:
     lit_config.note('using auto diff lib: %r' % config.cladlib)
 
 lit.util.usePlatformSdkOnDarwin(config, lit_config)
@@ -325,3 +324,6 @@ if libcudart_path is not None:
         config.substitutions.append(('%cudasmlevel', '--cuda-gpu-arch=sm_' + os.environ['CLING_TEST_CUDA_SM_LEVEL']))
     else:
         config.substitutions.append(('%cudasmlevel', ""))
+
+if(config.have_enzyme):
+    config.available_features.add('Enzyme')
diff --git a/test/lit.site.cfg.in b/test/lit.site.cfg.in
@@ -9,6 +9,7 @@ config.llvm_libs_dir = "@LLVM_LIBS_DIR@"
 config.clad_obj_root = "@CLAD_BINARY_DIR@"
 config.target_triple = "@TARGET_TRIPLE@"
 config.shlibext = "@TARGET_SHLIBEXT@"
+config.have_enzyme = "@ENABLE_ENZYME_BACKEND@"
 
 # Support substitution of the tools and libs dirs with user parameters. This is
 # used when we can't determine the tool dir at configuration time.