[llvm] [IndVarSimplify] Fix Masking Issue by Adding nsw/nuw Flags to Trunc Instruction (PR #150179)

[Luke Lau via llvm-commits]

================
@@ -1049,9 +1049,28 @@ linearFunctionTestReplace(Loop *L, BasicBlock *ExitingBB,
     if (Extended) {
       bool Discard;
       L->makeLoopInvariant(ExitCnt, Discard);
-    } else
+    } else{
       CmpIndVar = Builder.CreateTrunc(CmpIndVar, ExitCnt->getType(),
                                       "lftr.wideiv");
+
+      // Set the correct wrap flag to avoid the masking issue.
+      Instruction *TruncInst = dyn_cast<Instruction>(CmpIndVar);
+      
+      // The TruncatedIV is incrementing.
+      if (const SCEVAddRecExpr *TruncAR =
+              dyn_cast<SCEVAddRecExpr>(TruncatedIV)) {
+        // If TruncIV does not cause self-wrap, explicitly add the nsw and nuw
+        // flags to TruncInst.
+        if (TruncAR->hasNoSelfWrap()) {
+          TruncInst->setHasNoSignedWrap();
+          TruncInst->setHasNoUnsignedWrap();
----------------
lukel97 wrote:

> Given that TruncIV is marked as nw and truncation does not change the sign, can we assume that the nw flag of TruncIV retains the sign of the original IV? Specifically, if the original IV is signed, does TruncIV's nw imply that it behaves like nsw?

I think in general truncation can change the sign, e.g. i16 0x8000 truncated to i8 will be 0x00, so the sign can change. 

I tried out just adding the nuw/nsw wrap flags if the condition predicate is unsigned or signed, but it looks like that's unsound:

```diff
--- a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
+++ b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -1050,8 +1050,10 @@ linearFunctionTestReplace(Loop *L, BasicBlock *ExitingBB,
       bool Discard;
       L->makeLoopInvariant(ExitCnt, Discard);
     } else
-      CmpIndVar = Builder.CreateTrunc(CmpIndVar, ExitCnt->getType(),
-                                      "lftr.wideiv");
+      CmpIndVar =
+          Builder.CreateTrunc(CmpIndVar, ExitCnt->getType(), "lftr.wideiv",
+                              cast<ICmpInst>(BI->getCondition())->isUnsigned(),
+                              cast<ICmpInst>(BI->getCondition())->isSigned());
   }
diff --git a/llvm/test/Transforms/IndVarSimplify/lftr.ll b/llvm/test/Transforms/IndVarSimplify/lftr.ll
index 5ee62ba357ab..3825a49563a6 100644
--- a/llvm/test/Transforms/IndVarSimplify/lftr.ll
+++ b/llvm/test/Transforms/IndVarSimplify/lftr.ll
@@ -415,7 +415,7 @@ define void @wide_trip_count_test1(ptr %autoc,
 ; CHECK-NEXT:    [[ADD3:%.*]] = fadd float [[TEMP2]], [[MUL]]
 ; CHECK-NEXT:    store float [[ADD3]], ptr [[ARRAYIDX2]], align 4
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nuw nsw i64 [[INDVARS_IV]], 1
-; CHECK-NEXT:    [[LFTR_WIDEIV:%.*]] = trunc i64 [[INDVARS_IV_NEXT]] to i32
+; CHECK-NEXT:    [[LFTR_WIDEIV:%.*]] = trunc nuw i64 [[INDVARS_IV_NEXT]] to i32
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[LFTR_WIDEIV]], [[SUB]]
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_BODY]], label [[FOR_END_LOOPEXIT:%.*]]
 ; CHECK:       for.end.loopexit:
```

Passing the change to wide_trip_count_test1 through alive2 shows that this eventually triggers poison where it didn't before.

I wonder if it's possible to infer the NUW/NSW flags on TruncatedIV where possible? I think part of the NUW/NSW information is being lost when the checks are pulled to outside the loop, e.g. in your C example from the original issue:

```llvm
; Function Attrs: nounwind vscale_range(2,1024)
define dso_local void @func(ptr noundef captures(none) %result, i32 noundef signext %start) local_unnamed_addr #0 {
entry:
  %cmp3 = icmp slt i32 %start, 100
  br i1 %cmp3, label %for.body.preheader, label %for.cond.cleanup

for.body.preheader:                               ; preds = %entry
  %0 = sext i32 %start to i64
  br label %for.body

for.cond.cleanup.loopexit:                        ; preds = %for.body
  br label %for.cond.cleanup

for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
  ret void

for.body:                                         ; preds = %for.body.preheader, %for.body
  %indvars.iv = phi i64 [ %0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
  %i.04 = phi i32 [ %inc, %for.body ], [ %start, %for.body.preheader ]
  %idxprom = sext i32 %i.04 to i64
  %arrayidx = getelementptr inbounds i32, ptr %result, i64 %indvars.iv
  %1 = load i32, ptr %arrayidx, align 4, !tbaa !6
  %add = add nsw i32 %1, 1
  store i32 %add, ptr %arrayidx, align 4, !tbaa !6
  %indvars.iv.next = add nsw i64 %indvars.iv, 1
  %inc = add nsw i32 %i.04, 1
  %cmp = icmp slt i64 %indvars.iv, 99
  br i1 %cmp, label %for.body, label %for.cond.cleanup.loopexit, !llvm.loop !10
}
```

I don't think SCEV sees the `%cmp3 = icmp slt i32 %start, 100` condition and so it doesn't realise that %indvars.iv.next can't signed wrap?

https://github.com/llvm/llvm-project/pull/150179