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

[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();
----------------
buggfg 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 in i32?

We found the `linearFunctionTestReplace()`'s design can help determine whether truncation will change the sign:

1. The induction variable (`IV`) must be a `LoopCounter`, so its step is guaranteed to be `1`.
2. The `ICmpInst::Predicate` can only be `eq` or `ne`,  which means that `ExitCnt` must be the final value of `IV`.

Therefore, when the initial value `Start` of `IV` does not exceed `ExitCntSize`, the range `[start, end)` of `IV` will not cause signed or unsigned wrap.

For instance, in `llvm/test/Transforms/IndVarSimplify/lftr.ll`, the initial value `start = 68719476736 = 2^9` is within the range of `i32`, the step is `1`, and the final value `%sub` is also `i32`. Thus, the `IV` remains within the range of i32. `temp3 = trunc nuw i64 %indvars.iv.next to i32` is valid. 

So I propose to add a check for the type width of the initial `IV` value. If it matches the target type for truncation, then annotate the Trunc Instruction with `nsw` or `nuw` flag. Looking forward to your reply :)

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