Add gc_loaded support (#2067)

* Add gc_loaded support * Fix gc loaded * Fix fmt
EnzymeAD · Sep 2, 2024 · 457320d · 457320d
1 parent c810a63
commit 457320d
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 4 deletions.
diff --git a/.github/workflows/enzyme-mlir.yml b/.github/workflows/enzyme-mlir.yml
@@ -55,13 +55,11 @@ jobs:
       if: steps.cache-mlir.outputs.cache-hit != 'true'
       working-directory: 'llvm-project'
       run: |
-        mkdir mlir-build && cd mlir-build
-        cmake ../llvm -GNinja -DLLVM_ENABLE_PROJECTS="llvm;clang;mlir;openmp" -DCMAKE_BUILD_TYPE=${{ matrix.llbuild }} \
+        mkdir mlir-build && cd mlir-build && cmake ../llvm -GNinja -DLLVM_ENABLE_PROJECTS="llvm;clang;mlir;openmp" -DCMAKE_BUILD_TYPE=${{ matrix.llbuild }} \
         -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_TARGETS_TO_BUILD=X86 \
         -DLLVM_USE_LINKER=gold -DLLVM_PARALLEL_LINK_JOBS=2 \
         -DCLANG_ENABLE_STATIC_ANALYZER=OFF -DCLANG_ENABLE_ARCMT=OFF \
-        -DLLVM_OPTIMIZED_TABLEGEN=ON
-        ninja
+        -DLLVM_OPTIMIZED_TABLEGEN=ON && ninja || echo "already built"
 
     - name: Enzyme build
       working-directory: 'Enzyme'

diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -623,6 +623,9 @@ bool ActivityAnalyzer::isFunctionArgumentConstant(CallInst *CI, Value *val) {
   if (Name == "MPI_Waitall" || Name == "PMPI_Waitall")
     return val != CI->getOperand(1);
 
+  if (Name == "julia.gc_loaded")
+    return val != CI->getOperand(1);
+
   // TODO interprocedural detection
   // Before potential introprocedural detection, any function without definition
   // may to be assumed to have an active use
@@ -654,6 +657,11 @@ static inline void propagateArgumentInformation(
     return;
   }
 
+  if (Name == "julia.gc_loaded") {
+    propagateFromOperand(CI.getArgOperand(0));
+    return;
+  }
+
   if (Name == "julia.call" || Name == "julia.call2") {
 #if LLVM_VERSION_MAJOR >= 14
     for (size_t i = 1; i < CI.arg_size(); i++)

diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
@@ -3622,6 +3622,43 @@ bool AdjointGenerator::handleKnownCallDerivatives(
     return true;
   }
 
+  if (funcName == "julia.gc_loaded") {
+    if (gutils->isConstantValue(&call)) {
+      eraseIfUnused(call);
+      return true;
+    }
+    auto ifound = gutils->invertedPointers.find(&call);
+    assert(ifound != gutils->invertedPointers.end());
+
+    auto placeholder = cast<PHINode>(&*ifound->second);
+
+    bool needShadow =
+        DifferentialUseAnalysis::is_value_needed_in_reverse<QueryType::Shadow>(
+            gutils, &call, Mode, oldUnreachable);
+    if (!needShadow) {
+      gutils->invertedPointers.erase(ifound);
+      gutils->erase(placeholder);
+      eraseIfUnused(call);
+      return true;
+    }
+
+    Value *ptr0shadow = gutils->invertPointerM(call.getArgOperand(0), BuilderZ);
+    Value *ptr1shadow = gutils->invertPointerM(call.getArgOperand(1), BuilderZ);
+
+    Value *val = applyChainRule(
+        call.getType(), BuilderZ,
+        [&](Value *v1, Value *v2) -> Value * {
+          Value *args[2] = {v1, v2};
+          return BuilderZ.CreateCall(called, args);
+        },
+        ptr0shadow, ptr1shadow);
+
+    gutils->replaceAWithB(placeholder, val);
+    gutils->erase(placeholder);
+    eraseIfUnused(call);
+    return true;
+  }
+
   if (funcName == "julia.pointer_from_objref") {
     if (gutils->isConstantValue(&call)) {
       eraseIfUnused(call);