[llvm] [CVP] Freeze Y when expanding urem x, y with X < 2Y (PR #84390)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/84390

We're going from a single use to two independent uses, we need these two to see consistent values for undef.  As an example, consider x = 0x2 when y = 0b00u1.  If the sub use picks 0b0001 and the cmp use picks 0b0011, that would be incorrect.

>From 02338de2393603e81fcd4edfc4142ad106b3fb39 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 7 Mar 2024 14:09:04 -0800
Subject: [PATCH] [CVP] Freeze Y when expanding urem x, y with X < 2Y

We're going from a single use to two independent uses, we need these
two to see consistent values for undef.  As an example, consider
x = 0x2 when y = 0b00u1.  If the sub use picks 0b0001 and the cmp use
picks 0b0011, that would be incorrect.
---
 .../Scalar/CorrelatedValuePropagation.cpp     |  9 ++++++---
 .../urem-expansion.ll                         | 20 +++++++++++--------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
index 7a2011888ab008..de3bfb57b538d3 100644
--- a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
+++ b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
@@ -805,9 +805,12 @@ static bool expandUDivOrURem(BinaryOperator *Instr, const ConstantRange &XCR,
     Value *FrozenX = X;
     if (!isGuaranteedNotToBeUndef(X))
       FrozenX = B.CreateFreeze(X, X->getName() + ".frozen");
-    auto *AdjX = B.CreateNUWSub(FrozenX, Y, Instr->getName() + ".urem");
-    auto *Cmp =
-        B.CreateICmp(ICmpInst::ICMP_ULT, FrozenX, Y, Instr->getName() + ".cmp");
+    Value *FrozenY = Y;
+    if (!isGuaranteedNotToBeUndef(Y))
+      FrozenY = B.CreateFreeze(Y, Y->getName() + ".frozen");
+    auto *AdjX = B.CreateNUWSub(FrozenX, FrozenY, Instr->getName() + ".urem");
+    auto *Cmp = B.CreateICmp(ICmpInst::ICMP_ULT, FrozenX, FrozenY,
+                             Instr->getName() + ".cmp");
     ExpandedOp = B.CreateSelect(Cmp, FrozenX, AdjX);
   } else {
     auto *Cmp =
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/urem-expansion.ll b/llvm/test/Transforms/CorrelatedValuePropagation/urem-expansion.ll
index 2af8c8f23bbd3e..8e276d010fdd1b 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/urem-expansion.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/urem-expansion.ll
@@ -198,8 +198,9 @@ define i8 @variable.v4(i8 %x, i8 %y) {
 ; CHECK-NEXT:    [[CMP_Y_UPPER:%.*]] = icmp ule i8 [[Y]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP_Y_UPPER]])
 ; CHECK-NEXT:    [[X_FROZEN:%.*]] = freeze i8 [[X]]
-; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y]]
-; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y]]
+; CHECK-NEXT:    [[Y_FROZEN:%.*]] = freeze i8 [[Y]]
+; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y_FROZEN]]
+; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[REM:%.*]] = select i1 [[REM_CMP]], i8 [[X_FROZEN]], i8 [[REM_UREM]]
 ; CHECK-NEXT:    ret i8 [[REM]]
 ;
@@ -217,8 +218,9 @@ define i8 @variable.v4.range(ptr %x.ptr, ptr %y.ptr) {
 ; CHECK-NEXT:    [[X:%.*]] = load i8, ptr [[X_PTR:%.*]], align 1, !range [[RNG0]]
 ; CHECK-NEXT:    [[Y:%.*]] = load i8, ptr [[Y_PTR:%.*]], align 1, !range [[RNG1:![0-9]+]]
 ; CHECK-NEXT:    [[X_FROZEN:%.*]] = freeze i8 [[X]]
-; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y]]
-; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y]]
+; CHECK-NEXT:    [[Y_FROZEN:%.*]] = freeze i8 [[Y]]
+; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y_FROZEN]]
+; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[REM:%.*]] = select i1 [[REM_CMP]], i8 [[X_FROZEN]], i8 [[REM_UREM]]
 ; CHECK-NEXT:    ret i8 [[REM]]
 ;
@@ -236,8 +238,9 @@ define i8 @variable.v5(i8 %x, i8 %y) {
 ; CHECK-NEXT:    [[CMP_Y_UPPER:%.*]] = icmp ule i8 [[Y]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP_Y_UPPER]])
 ; CHECK-NEXT:    [[X_FROZEN:%.*]] = freeze i8 [[X]]
-; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y]]
-; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y]]
+; CHECK-NEXT:    [[Y_FROZEN:%.*]] = freeze i8 [[Y]]
+; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y_FROZEN]]
+; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[REM:%.*]] = select i1 [[REM_CMP]], i8 [[X_FROZEN]], i8 [[REM_UREM]]
 ; CHECK-NEXT:    ret i8 [[REM]]
 ;
@@ -259,8 +262,9 @@ define i8 @variable.v6(i8 %x, i8 %y) {
 ; CHECK-NEXT:    [[CMP_Y_UPPER:%.*]] = icmp ule i8 [[Y]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP_Y_UPPER]])
 ; CHECK-NEXT:    [[X_FROZEN:%.*]] = freeze i8 [[X]]
-; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y]]
-; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y]]
+; CHECK-NEXT:    [[Y_FROZEN:%.*]] = freeze i8 [[Y]]
+; CHECK-NEXT:    [[REM_UREM:%.*]] = sub nuw i8 [[X_FROZEN]], [[Y_FROZEN]]
+; CHECK-NEXT:    [[REM_CMP:%.*]] = icmp ult i8 [[X_FROZEN]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[REM:%.*]] = select i1 [[REM_CMP]], i8 [[X_FROZEN]], i8 [[REM_UREM]]
 ; CHECK-NEXT:    ret i8 [[REM]]
 ;