slice activity analysis

EnzymeAD · Dec 28, 2024 · bac0f7e · bac0f7e
1 parent eeb6200
commit bac0f7e
Showing 1 changed file with 35 additions and 0 deletions.
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -579,6 +579,11 @@ 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 +665,13 @@ 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));
@@ -1601,6 +1613,29 @@ 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() &&