Improve DifferentialUseAnalysis (rust-lang#709)

* Improve DifferentialUseAnalysis

* Fix test

Co-authored-by: William S. Moses <[email protected]>

Author: tgymnich

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -107,8 +107,6 @@ static inline bool is_use_directly_needed_in_reverse(
 
   if (isa<CmpInst>(user) || isa<BranchInst>(user) || isa<ReturnInst>(user) ||
       isa<FPExtInst>(user) || isa<FPTruncInst>(user) ||
-      (isa<InsertElementInst>(user) &&
-       cast<InsertElementInst>(user)->getOperand(2) != val) ||
       (isa<ExtractElementInst>(user) &&
        cast<ExtractElementInst>(user)->getIndexOperand() != val)
 #if LLVM_VERSION_MAJOR >= 10
@@ -121,6 +119,44 @@ static inline bool is_use_directly_needed_in_reverse(
     return false;
   }
 
+  if (auto IEI = dyn_cast<InsertElementInst>(user)) {
+    // Only need the index in the reverse, so if the value is not
+    // the index, short circuit and say we don't need
+    if (IEI->getOperand(2) != val) {
+      return false;
+    }
+    // The index is only needed in the reverse if the value being inserted
+    // is a possible active floating point value
+    if (gutils->isConstantValue(const_cast<InsertElementInst *>(IEI)) ||
+        TR.query(const_cast<InsertElementInst *>(IEI))[{-1}] ==
+            BaseType::Pointer)
+      return false;
+    // Otherwise, we need the value.
+    return true;
+  }
+
+  if (auto IVI = dyn_cast<InsertValueInst>(user)) {
+    // Only need the index in the reverse, so if the value is not
+    // the index, short circuit and say we don't need
+    bool valueIsIndex = false;
+    for (unsigned i = 2; i < IVI->getNumOperands(); ++i) {
+      if (IVI->getOperand(i) == val) {
+        valueIsIndex = true;
+      }
+    }
+
+    if (!valueIsIndex)
+      return false;
+
+    // The index is only needed in the reverse if the value being inserted
+    // is a possible active floating point value
+    if (gutils->isConstantValue(const_cast<InsertValueInst *>(IVI)) ||
+        TR.query(const_cast<InsertValueInst *>(IVI))[{-1}] == BaseType::Pointer)
+      return false;
+    // Otherwise, we need the value.
+    return true;
+  }
+
   Intrinsic::ID ID = Intrinsic::not_intrinsic;
   if (auto II = dyn_cast<IntrinsicInst>(user)) {
     ID = II->getIntrinsicID();

enzyme/test/Enzyme/ForwardModeVector/vecsquare.ll

@@ -56,32 +56,26 @@ entry:
 ; CHECK-NEXT:    [[TMP26:%.*]] = insertvalue [2 x <4 x float>] undef, <4 x float> %"id'ipsv", 0
 ; CHECK-NEXT:    %"id'ipsv1" = shufflevector <4 x float> [[TMP10]], <4 x float> [[TMP22]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
 ; CHECK-NEXT:    [[TMP29:%.*]] = insertvalue [2 x <4 x float>] [[TMP26]], <4 x float> %"id'ipsv1", 1
-; CHECK-NEXT:    [[ID:%.*]] = shufflevector <4 x float> [[SQ]], <4 x float> [[CB]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
 ; CHECK-NEXT:    %"res1'ipee" = extractelement <4 x float> %"id'ipsv", i32 1
 ; CHECK-NEXT:    [[TMP31:%.*]] = insertvalue [2 x float] undef, float %"res1'ipee", 0
 ; CHECK-NEXT:    %"res1'ipee2" = extractelement <4 x float> %"id'ipsv1", i32 1
 ; CHECK-NEXT:    [[TMP33:%.*]] = insertvalue [2 x float] [[TMP31]], float %"res1'ipee2", 1
-; CHECK-NEXT:    [[RES1:%.*]] = extractelement <4 x float> [[ID]], i32 1
 ; CHECK-NEXT:    %"res2'ipee" = extractelement <4 x float> %"id'ipsv", i32 2
 ; CHECK-NEXT:    [[TMP35:%.*]] = insertvalue [2 x float] undef, float %"res2'ipee", 0
 ; CHECK-NEXT:    %"res2'ipee3" = extractelement <4 x float> %"id'ipsv1", i32 2
 ; CHECK-NEXT:    [[TMP37:%.*]] = insertvalue [2 x float] [[TMP35]], float %"res2'ipee3", 1
-; CHECK-NEXT:    [[RES2:%.*]] = extractelement <4 x float> [[ID]], i32 2
 ; CHECK-NEXT:    %"res3'ipee" = extractelement <4 x float> %"id'ipsv", i32 3
 ; CHECK-NEXT:    [[TMP39:%.*]] = insertvalue [2 x float] undef, float %"res3'ipee", 0
 ; CHECK-NEXT:    %"res3'ipee4" = extractelement <4 x float> %"id'ipsv1", i32 3
 ; CHECK-NEXT:    [[TMP41:%.*]] = insertvalue [2 x float] [[TMP39]], float %"res3'ipee4", 1
-; CHECK-NEXT:    [[RES3:%.*]] = extractelement <4 x float> [[ID]], i32 3
 ; CHECK-NEXT:    %"agg1'ipiv" = insertvalue { float, float, float } zeroinitializer, float %"res1'ipee", 0
 ; CHECK-NEXT:    [[TMP43:%.*]] = insertvalue [2 x { float, float, float }] undef, { float, float, float } %"agg1'ipiv", 0
 ; CHECK-NEXT:    %"agg1'ipiv5" = insertvalue { float, float, float } zeroinitializer, float %"res1'ipee2", 0
 ; CHECK-NEXT:    [[TMP45:%.*]] = insertvalue [2 x { float, float, float }] [[TMP43]], { float, float, float } %"agg1'ipiv5", 1
-; CHECK-NEXT:    [[AGG1:%.*]] = insertvalue { float, float, float } undef, float [[RES1]], 0
 ; CHECK-NEXT:    %"agg2'ipiv" = insertvalue { float, float, float } %"agg1'ipiv", float %"res2'ipee", 1
 ; CHECK-NEXT:    [[TMP48:%.*]] = insertvalue [2 x { float, float, float }] undef, { float, float, float } %"agg2'ipiv", 0
 ; CHECK-NEXT:    %"agg2'ipiv6" = insertvalue { float, float, float } %"agg1'ipiv5", float %"res2'ipee3", 1
 ; CHECK-NEXT:    [[TMP51:%.*]] = insertvalue [2 x { float, float, float }] [[TMP48]], { float, float, float } %"agg2'ipiv6", 1
-; CHECK-NEXT:    [[AGG2:%.*]] = insertvalue { float, float, float } [[AGG1]], float [[RES2]], 1
 ; CHECK-NEXT:    %"agg3'ipiv" = insertvalue { float, float, float } %"agg2'ipiv", float %"res3'ipee", 2
 ; CHECK-NEXT:    [[TMP54:%.*]] = insertvalue [2 x { float, float, float }] undef, { float, float, float } %"agg3'ipiv", 0
 ; CHECK-NEXT:    %"agg3'ipiv7" = insertvalue { float, float, float } %"agg2'ipiv6", float %"res3'ipee4", 2