slice activity analysis

EnzymeAD · Dec 28, 2024 · 7e4b13f · 7e4b13f
1 parent 9fe27b8
commit 7e4b13f
Showing 1 changed file with 31 additions and 0 deletions.
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -579,6 +579,10 @@ bool ActivityAnalyzer::isFunctionArgumentConstant(CallInst *CI, Value *val) {
   if (Name == "jl_reshape_array" || Name == "ijl_reshape_array")
     return val != CI->getArgOperand(1);
 
+  // Only the 0-th arg impacts activity
+  if (Name == "jl_genericmemory_copy_slice" || Name == "ijl_genericmemory_copy_slice")
+    return val != CI->getArgOperand(0);
+
   // Allocations, deallocations, and c++ guards don't impact the activity
   // of arguments
   if (isAllocationFunction(Name, TLI) || isDeallocationFunction(Name, TLI))
@@ -660,6 +664,12 @@ static inline void propagateArgumentInformation(
     return;
   }
 
+  // Only the 0-th arg impacts activity
+  if (Name == "jl_genericmemory_copy_slice" || Name == "ijl_genericmemory_copy_slice") {
+    propagateFromOperand(CI.getArgOperand(1));
+    return;
+  }
+
   // Only the 1-th arg impacts activity
   if (Name == "jl_reshape_array" || Name == "ijl_reshape_array") {
     propagateFromOperand(CI.getArgOperand(1));
@@ -1583,6 +1593,27 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
           return true;
         }
 
+        if (funcName == "jl_genericmemory_copy_slice" || funcName == "ijl_genericmemory_copy_slice") {
+          if (directions == UP) {
+            if (isConstantValue(TR, op->getArgOperand(0))) {
+              InsertConstantValue(TR, Val);
+              return true;
+            }
+          } else {
+            if (UpHypothesis->isConstantValue(TR, op->getArgOperand(0))) {
+              InsertConstantValue(TR, Val);
+              insertConstantsFrom(TR, *UpHypothesis);
+              return true;
+            }
+          }
+          if (EnzymeEnableRecursiveHypotheses) {
+            ReEvaluateValueIfInactiveValue[op->getArgOperand(0)].insert(Val);
+            if (TmpOrig != Val) {
+              ReEvaluateValueIfInactiveValue[op->getArgOperand(0)].insert(TmpOrig);
+            }
+          }
+        }
+
         // If requesting empty unknown functions to be considered inactive,
         // abide by those rules
         if (called && EnzymeEmptyFnInactive && called->empty() &&