Enable get return diffe type to be used in more scenarios (#1300)

* Enable get return diffe type to be used in more scenarios * Fix subretused of gc allocation
EnzymeAD · Jun 26, 2023 · bccd9ff · bccd9ff
1 parent 366d3c0
commit bccd9ff
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -9618,7 +9618,7 @@ class AdjointGenerator
  funcName == "ijl_alloc_array_3d" || funcName == "ijl_array_copy" ||
  funcName == "julia.gc_alloc_obj" || funcName == "jl_gc_alloc_typed" ||
  funcName == "ijl_gc_alloc_typed") {
- if (unnecessaryValues.count(&call)) {
+ if (!subretused) {
  eraseIfUnused(call, /*erase*/ true, /*check*/ false);
  return;
  }

diff --git a/enzyme/Enzyme/CApi.cpp b/enzyme/Enzyme/CApi.cpp
@@ -397,8 +397,8 @@ EnzymeGradientUtilsGetReturnDiffeType(GradientUtils *G, LLVMValueRef oval,
  uint8_t *needsShadow) {
  bool needsPrimalB;
  bool needsShadowB;
- auto res = (CDIFFE_TYPE)(G->getReturnDiffeType(cast<CallInst>(unwrap(oval)),
- &needsPrimalB, &needsShadowB));
+ auto res = (CDIFFE_TYPE)(G->getReturnDiffeType(unwrap(oval), &needsPrimalB,
+ &needsShadowB));
  if (needsPrimal)
  *needsPrimal = needsPrimalB;
  if (needsShadow)

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -4339,7 +4339,7 @@ GradientUtils *GradientUtils::CreateFromClone(
  return res;
 }
 
-DIFFE_TYPE GradientUtils::getReturnDiffeType(llvm::CallInst *orig,
+DIFFE_TYPE GradientUtils::getReturnDiffeType(llvm::Value *orig,
  bool *primalReturnUsedP,
  bool *shadowReturnUsedP) const {
  bool shadowReturnUsed = false;

diff --git a/enzyme/Enzyme/GradientUtils.h b/enzyme/Enzyme/GradientUtils.h
@@ -376,7 +376,7 @@ class GradientUtils : public CacheUtility {
 public:
  DIFFE_TYPE getDiffeType(llvm::Value *v, bool foreignFunction) const;
 
- DIFFE_TYPE getReturnDiffeType(llvm::CallInst *orig, bool *primalReturnUsedP,
+ DIFFE_TYPE getReturnDiffeType(llvm::Value *orig, bool *primalReturnUsedP,
  bool *shadowReturnUsedP) const;
 
  static GradientUtils *