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

[via llvm-commits]

https://github.com/vfdff created https://github.com/llvm/llvm-project/pull/74021

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

Fix https://github.com/llvm/llvm-project/issues/71782

>From 833ded650766d021343b8a488d3b48b02fd1d7dc Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Thu, 9 Nov 2023 02:42:46 -0500
Subject: [PATCH] [ValueTracking] Support srem/urem for
 isKnownNonNullFromDominatingCondition

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

Fix https://github.com/llvm/llvm-project/issues/71782
---
 llvm/lib/Analysis/ValueTracking.cpp           |  4 +-
 .../ValueTracking/select-known-non-zero.ll    | 56 ++++++++++++++++++-
 .../Transforms/InstCombine/zext-or-icmp.ll    |  4 +-
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 250ce739ea5147c..eae5a324ce25dd3 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2186,7 +2186,9 @@ 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_URem(m_Value(), m_Specific(V))) ||
+         match(U, m_SRem(m_Value(), m_Specific(V)))) &&
         isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
       return true;
 
diff --git a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll b/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
index 1dc88412041d34f..53ed4485c94f086 100644
--- a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
+++ b/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
+}
diff --git a/llvm/test/Transforms/InstCombine/zext-or-icmp.ll b/llvm/test/Transforms/InstCombine/zext-or-icmp.ll
index bc0e4bdce29b595..9ec3ddc80c57f7f 100644
--- a/llvm/test/Transforms/InstCombine/zext-or-icmp.ll
+++ b/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]]