[InstCombine] handle (X + A) % Op1 for small X, A

Author: DesWurstes

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

@@ -2461,6 +2461,24 @@ Instruction *InstCombinerImpl::visitURem(BinaryOperator &I) {
     }
   }
 
+  // For "(X + A) % Op1" and if (X u< Op1 && A u< Op1)
+  // => (X + A) >= Op1 ? X + A - Op1 : X + A .
+  Value *A = nullptr;
+  if (match(Op0, m_Add(m_Value(X), m_Value(A)))) {
+    Value *Val_X =
+        simplifyICmpInst(ICmpInst::ICMP_ULT, X, Op1, SQ.getWithInstruction(&I));
+    Value *Val_A =
+        simplifyICmpInst(ICmpInst::ICMP_ULT, A, Op1, SQ.getWithInstruction(&I));
+    if (Val_X && match(Val_X, m_One()) && Val_A && match(Val_A, m_One())) {
+      Value *FrozenOp0 = Op0;
+      if (!isGuaranteedNotToBeUndef(Op0))
+        FrozenOp0 = Builder.CreateFreeze(Op0, Op0->getName() + ".frozen");
+      Value *Cmp = Builder.CreateICmpUGE(FrozenOp0, Op1);
+      Value *Sub = Builder.CreateSub(FrozenOp0, Op1);
+      return SelectInst::Create(Cmp, Sub, FrozenOp0);
+    }
+  }
+
   return nullptr;
 }
 

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

@@ -27,7 +27,10 @@ define i8 @urem_assume_a(i8 %x, i8 %n, i8 %a) {
 ; CHECK-NEXT:    [[CMP_A:%.*]] = icmp ult i8 [[A:%.*]], [[N]]
 ; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP_A]])
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X_FR]], [[A]]
-; CHECK-NEXT:    [[OUT:%.*]] = urem i8 [[ADD]], [[N]]
+; CHECK-NEXT:    [[ADD_FROZEN:%.*]] = freeze i8 [[ADD]]
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp ult i8 [[ADD_FROZEN]], [[N]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[DOTNOT]], i8 0, i8 [[N]]
+; CHECK-NEXT:    [[OUT:%.*]] = sub i8 [[ADD_FROZEN]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[OUT]]
 ;
   %cmp = icmp ult i8 %x, %n
@@ -231,7 +234,10 @@ define i8 @urem_without_assume_a(i8 %arg, i8 %arg2, i8 %a) {
 ; CHECK-NEXT:    [[X_REM:%.*]] = urem i8 [[ARG:%.*]], [[ARG2:%.*]]
 ; CHECK-NEXT:    [[A_REM:%.*]] = urem i8 [[A:%.*]], [[ARG2]]
 ; CHECK-NEXT:    [[ADDD:%.*]] = add i8 [[X_REM]], [[A_REM]]
-; CHECK-NEXT:    [[OUT:%.*]] = urem i8 [[ADDD]], [[ARG2]]
+; CHECK-NEXT:    [[ADD_FROZEN:%.*]] = freeze i8 [[ADDD]]
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp ult i8 [[ADD_FROZEN]], [[ARG2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[DOTNOT]], i8 0, i8 [[ARG2]]
+; CHECK-NEXT:    [[OUT:%.*]] = sub i8 [[ADD_FROZEN]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[OUT]]
 ;
   %x_rem = urem i8 %arg, %arg2