[DAGCombiner] Fold add (mul x, C), x to mul x, C+1

[nikic]

Summary

While this is normally non-canonical IR, this pattern can appear during SDAG lowering if the add is actually a getelementptr, as illustrated in @test_ptr. This pattern comes up when doing provenance-aware high-bit pointer tagging.

Proof: https://alive2.llvm.org/ce/z/DLoEcs

[RKSimon]

LGTM with one minor

llvm/test/CodeGen/X86/add-of-mul.ll

24

pre-add this new test.

[goldstein.w.n]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3061

Does the equivilent ever come up for sub?

[nikic]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3061

I don't think so. You'd need something like gep(x, sub(0, mul(x, C))) for that, and that would be canonicalized to gep(x, mul(x, -C)) rather than becoming sub(x, mul(x, C)) in the backend.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (N0.getOpcode() == ISD::SUB && N0.hasOneUse()) {
     // Hoist one-use subtraction from non-opaque constant:
     //   (C - x) + y  ->  (y - x) + C
     if (isConstantOrConstantVector(N0.getOperand(0), /*NoOpaques=*/true)) {
       SDValue Sub = DAG.getNode(ISD::SUB, DL, VT, N1, N0.getOperand(1));
       return DAG.getNode(ISD::ADD, DL, VT, Sub, N0.getOperand(0));
     }
   }
 
+  // add (mul x, C), x -> mul x, C+1
+  if (N0.getOpcode() == ISD::MUL && N0.getOperand(0) == N1 &&
+      isConstantOrConstantVector(N0.getOperand(1), /*NoOpaques=*/true) &&
+      N0.hasOneUse()) {
+    SDValue NewC = DAG.getNode(ISD::ADD, DL, VT, N0.getOperand(1),
+                               DAG.getConstant(1, DL, VT));
+    return DAG.getNode(ISD::MUL, DL, VT, N0.getOperand(0), NewC);
+  }
+
   // If the target's bool is represented as 0/1, prefer to make this 'sub 0/1'
   // rather than 'add 0/-1' (the zext should get folded).
   // add (sext i1 Y), X --> sub X, (zext i1 Y)
   if (N0.getOpcode() == ISD::SIGN_EXTEND &&
       N0.getOperand(0).getScalarValueSizeInBits() == 1 &&
       TLI.getBooleanContents(VT) == TargetLowering::ZeroOrOneBooleanContent) {
     SDValue ZExt = DAG.getNode(ISD::ZERO_EXTEND, DL, VT, N0.getOperand(0));
     return DAG.getNode(ISD::SUB, DL, VT, N1, ZExt);

llvm/test/CodeGen/AArch64/arm64-promote-const.ll

 define <16 x i8> @test2(<16 x i8> %arg) {
 entry:
 ; PROMOTED-LABEL: test2:
 ; In stress mode, constant vector are promoted
 ; PROMOTED: adrp [[PAGEADDR:x[0-9]+]], [[CSTV1:__PromotedConst.[0-9]+]]@PAGE
 ; PROMOTED: ldr q[[REGNUM:[0-9]+]], [[[PAGEADDR]], [[CSTV1]]@PAGEOFF]
 ; Destination register is defined by ABI
 ; PROMOTED-NEXT: add.16b v0, v0, v[[REGNUM]]
-; PROMOTED-NEXT: mla.16b v0, v0, v[[REGNUM]]
+; PROMOTED-NEXT: mls.16b v0, v0, v[[REGNUM]]
 ; PROMOTED-NEXT: ret
 
 ; REGULAR-LABEL: test2:
 ; Regular access is strickly the same as promoted access.
 ; The difference is that the address (and thus the space in memory) is not
 ; shared between constants
 ; REGULAR: adrp [[PAGEADDR:x[0-9]+]], [[CSTLABEL:lCP.*]]@PAGE
 ; REGULAR: ldr q[[REGNUM:[0-9]+]], [[[PAGEADDR]], [[CSTLABEL]]@PAGEOFF]
 ; Destination register is defined by ABI
 ; REGULAR-NEXT: add.16b v0, v0, v[[REGNUM]]
-; REGULAR-NEXT: mla.16b v0, v0, v[[REGNUM]]
+; REGULAR-NEXT: mls.16b v0, v0, v[[REGNUM]]
 ; REGULAR-NEXT: ret
   %add.i = add <16 x i8> %arg, <i8 -40, i8 -93, i8 -118, i8 -99, i8 -75, i8 -105, i8 74, i8 -110, i8 62, i8 -115, i8 -119, i8 -120, i8 34, i8 -124, i8 0, i8 -128>
   %mul.i = mul <16 x i8> %add.i, <i8 -40, i8 -93, i8 -118, i8 -99, i8 -75, i8 -105, i8 74, i8 -110, i8 62, i8 -115, i8 -119, i8 -120, i8 34, i8 -124, i8 0, i8 -128>
-  %add.i9 = add <16 x i8> %add.i, %mul.i
-  ret <16 x i8> %add.i9
+  %sub.i9 = sub <16 x i8> %add.i, %mul.i
+  ret <16 x i8> %sub.i9
 }
 
 ; Two different uses of the same constant in two different basic blocks,
 ; one dominates the other
 define <16 x i8> @test3(<16 x i8> %arg, i32 %path) {
 ; PROMOTED-LABEL: test3:
 ; In stress mode, constant vector are promoted
 ; Since, the constant is the same as the previous function,

llvm/test/CodeGen/X86/add-of-mul.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
 ; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
 
 define i32 @test_scalar(i32 %x) {
 ; CHECK-LABEL: test_scalar:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
-; CHECK-NEXT:    leal (%rdi,%rdi,2), %eax
-; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    leal (,%rdi,4), %eax
 ; CHECK-NEXT:    retq
   %mul = mul i32 %x, 3
   %add = add i32 %mul, %x
   ret i32 %add
 }
 
 define i32 @test_scalar_commuted(i32 %x) {
 ; CHECK-LABEL: test_scalar_commuted:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
-; CHECK-NEXT:    leal (%rdi,%rdi,2), %eax
-; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    leal (,%rdi,4), %eax
 ; CHECK-NEXT:    retq
   %mul = mul i32 %x, 3
   %add = add i32 %x, %mul
   ret i32 %add
 }
 
 define <4 x i32> @test_vector(<4 x i32> %x) {
 ; CHECK-LABEL: test_vector:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    paddd %xmm0, %xmm0
-; CHECK-NEXT:    paddd %xmm0, %xmm0
+; CHECK-NEXT:    pslld $2, %xmm0
 ; CHECK-NEXT:    retq
   %mul = mul <4 x i32> %x, <i32 3, i32 3, i32 3, i32 3>
   %add = add <4 x i32> %mul, %x
   ret <4 x i32> %add
 }
 
 define ptr @test_ptr(ptr %p) {
 ; CHECK-LABEL: test_ptr:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    leaq (%rdi,%rdi,2), %rax
-; CHECK-NEXT:    addq %rdi, %rax
+; CHECK-NEXT:    leaq (,%rdi,4), %rax
 ; CHECK-NEXT:    retq
   %addr = ptrtoint ptr %p to i64
   %mul = mul i64 %addr, 3
   %res = getelementptr i8, ptr %p, i64 %mul
   ret ptr %res
 }
 
 define i32 @test_extra_use(i32 %x, ptr %p) {