[ValueTracking] Support srem/urem for isKnownNonNullFromDominatingCondition (#74021)

Similar to div, the rem should also proof its second operand is
non-zero, otherwise it is a UB.

Fix #71782

Author: vfdff

llvm/lib/Analysis/ValueTracking.cpp

@@ -2186,7 +2186,8 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         return true;
     }
 
-    if (match(U, m_IDiv(m_Value(), m_Specific(V))) &&
+    if ((match(U, m_IDiv(m_Value(), m_Specific(V))) ||
+         match(U, m_IRem(m_Value(), m_Specific(V)))) &&
         isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
       return true;
 

llvm/test/Analysis/ValueTracking/select-known-non-zero.ll

@@ -2,6 +2,8 @@
 ; RUN: opt < %s -passes=instsimplify -S | FileCheck %s
 
 declare void @llvm.assume(i1)
+declare void @use(i64)
+declare void @use4(i4)
 
 define i1 @select_v_ne_fail(i8 %v, i8 %C, i8 %y) {
 ; CHECK-LABEL: @select_v_ne_fail(
@@ -446,4 +448,56 @@ define i64 @incorrect_safe_div_call_2(i64 %n, i64 %d) {
   ret i64 %3
 }
 
-declare void @use(i64)
+; https://alive2.llvm.org/ce/z/Si_B7b
+define i4 @icmp_urem(i4 %n, i4 %d) {
+; CHECK-LABEL: @icmp_urem(
+; CHECK-NEXT:    [[TMP1:%.*]] = urem i4 [[N:%.*]], [[D:%.*]]
+; CHECK-NEXT:    ret i4 [[TMP1]]
+;
+  %1 = icmp eq i4 %d, 0
+  %2 = urem i4 %n, %d
+  %3 = select i1 %1, i4 -1, i4 %2
+  ret i4 %3
+}
+
+define i4 @icmp_urem_clobber_by_call(i4 %n, i4 %d) {
+; CHECK-LABEL: @icmp_urem_clobber_by_call(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i4 [[D:%.*]], 0
+; CHECK-NEXT:    tail call void @use4(i4 [[D]])
+; CHECK-NEXT:    [[TMP2:%.*]] = urem i4 [[N:%.*]], [[D]]
+; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i4 -1, i4 [[TMP2]]
+; CHECK-NEXT:    ret i4 [[TMP3]]
+;
+  %1 = icmp eq i4 %d, 0
+  tail call void @use4(i4 %d)
+  %2 = urem i4 %n, %d
+  %3 = select i1 %1, i4 -1, i4 %2
+  ret i4 %3
+}
+
+; https://alive2.llvm.org/ce/z/Fn3Wac
+define i4 @icmp_srem(i4 %n, i4 %d) {
+; CHECK-LABEL: @icmp_srem(
+; CHECK-NEXT:    [[TMP1:%.*]] = srem i4 [[N:%.*]], [[D:%.*]]
+; CHECK-NEXT:    ret i4 [[TMP1]]
+;
+  %1 = icmp eq i4 %d, 0
+  %2 = srem i4 %n, %d
+  %3 = select i1 %1, i4 -1, i4 %2
+  ret i4 %3
+}
+
+define i4 @icmp_srem_clobber_by_call(i4 %n, i4 %d) {
+; CHECK-LABEL: @icmp_srem_clobber_by_call(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i4 [[D:%.*]], 0
+; CHECK-NEXT:    tail call void @use4(i4 [[D]])
+; CHECK-NEXT:    [[TMP2:%.*]] = srem i4 [[N:%.*]], [[D]]
+; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i4 -1, i4 [[TMP2]]
+; CHECK-NEXT:    ret i4 [[TMP3]]
+;
+  %1 = icmp eq i4 %d, 0
+  tail call void @use4(i4 %d)
+  %2 = srem i4 %n, %d
+  %3 = select i1 %1, i4 -1, i4 %2
+  ret i4 %3
+}

llvm/test/Transforms/InstCombine/zext-or-icmp.ll

@@ -231,9 +231,7 @@ define i1 @PR51762(ptr %i, i32 %t0, i16 %t1, ptr %p, ptr %d, ptr %f, i32 %p2, i1
 ; CHECK-NEXT:    [[INSERT_INSERT41:%.*]] = or i64 [[INSERT_SHIFT52]], [[INSERT_EXT39]]
 ; CHECK-NEXT:    [[REM:%.*]] = urem i64 [[S1]], [[INSERT_INSERT41]]
 ; CHECK-NEXT:    [[NE:%.*]] = icmp ne i64 [[REM]], 0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[INSERT_INSERT41]], 0
-; CHECK-NEXT:    [[SPEC_SELECT57:%.*]] = or i1 [[NE]], [[CMP]]
-; CHECK-NEXT:    [[LOR_EXT:%.*]] = zext i1 [[SPEC_SELECT57]] to i32
+; CHECK-NEXT:    [[LOR_EXT:%.*]] = zext i1 [[NE]] to i32
 ; CHECK-NEXT:    [[T2:%.*]] = load i32, ptr [[D:%.*]], align 4
 ; CHECK-NEXT:    [[CONV15:%.*]] = sext i16 [[T1]] to i32
 ; CHECK-NEXT:    [[CMP16:%.*]] = icmp sge i32 [[T2]], [[CONV15]]