[AArch64] Implement target hook function to decide folding (mul (add x, c1), c2)

[benshi001]

Summary

Prevent the folding in DAGCombine if it leads to worse code.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12152

Can this use AArch64TTIImpl::getIntImmCost in some way?

[benshi001]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12152

I think it is hard. Since the API

TargetTransformInfo
AArch64TargetMachine::getTargetTransformInfo(const Function &F) {
  return TargetTransformInfo(AArch64TTIImpl(this, F));
}

requires a Function object, but the target hook has only two SDValue parameters.

So the cheapest way is implementing a new but short isMOVZImm helper function.

[benshi001]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12152

I have simplified it to one additional helper function.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12152

It's more about not repeating the same (or similar) code. Can this at least use AArch64_IMM::expandMOVImm and check the number of instructions?

[benshi001]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12152

Thanks for your suggestion. Using AArch64_IMM::expandMOVImm makes the check more clear and accurate !!!

[dmgreen]

Thanks. LGTM

[kda]

Reverting.
Broke buildbot: https://lab.llvm.org/buildbot/#/builders/5/builds/11411

salient log snippet:

Testing:  0.. 10.. 20.. 30.. 40
FAIL: LLVM :: CodeGen/AArch64/urem-seteq-nonzero.ll (34455 of 78731)
******************** TEST 'LLVM :: CodeGen/AArch64/urem-seteq-nonzero.ll' FAILED ********************
Script:
--
: 'RUN: at line 2';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc -mtriple=aarch64-unknown-linux-gnu < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/AArch64/urem-seteq-nonzero.ll | /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/AArch64/urem-seteq-nonzero.ll
--
Exit Code: 2
Command Output (stderr):
--
/b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp:12208:58: runtime error: signed integer overflow: -2 * -6148914691236517205 cannot be represented in type 'long'
    #0 0x3967870 in llvm::AArch64TargetLowering::isMulAddWithConstProfitable(llvm::SDValue const&, llvm::SDValue const&) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp:12208:58
    #1 0x5cfdcde in (anonymous namespace)::DAGCombiner::isMulAddWithConstProfitable(llvm::SDNode*, llvm::SDValue&, llvm::SDValue&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:16879:11
    #2 0x5cb3dc3 in (anonymous namespace)::DAGCombiner::visitMUL(llvm::SDNode*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:3936:7
    #3 0x5caaafc in visit /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:1633:40
    #4 0x5caaafc in (anonymous namespace)::DAGCombiner::combine(llvm::SDNode*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:1748:10
    #5 0x5ca617e in Run /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:1556:18
    #6 0x5ca617e in llvm::SelectionDAG::Combine(llvm::CombineLevel, llvm::AAResults*, llvm::CodeGenOpt::Level) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp:23600:36
    #7 0x5e8d976 in llvm::SelectionDAGISel::CodeGenAndEmitDAG() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp:811:13
    #8 0x5e8d374 in llvm::SelectionDAGISel::SelectBasicBlock(llvm::ilist_iterator<llvm::ilist_detail::node_options<llvm::Instruction, true, false, void>, false, true>, llvm::ilist_iterator<llvm::ilist_detail::node_options<llvm::Instruction, true, false, void>, false, true>, bool&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp:727:3
    #9 0x5e8bc2f in llvm::SelectionDAGISel::SelectAllBasicBlocks(llvm::Function const&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp:1625:7
    #10 0x5e872ac in llvm::SelectionDAGISel::runOnMachineFunction(llvm::MachineFunction&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp:509:3
    #11 0x51b564e in llvm::MachineFunctionPass::runOnFunction(llvm::Function&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/CodeGen/MachineFunctionPass.cpp:72:13
    #12 0x570c8ea in llvm::FPPassManager::runOnFunction(llvm::Function&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1439:27
    #13 0x571544d in llvm::FPPassManager::runOnModule(llvm::Module&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1485:16
    #14 0x570d340 in runOnModule /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:1554:27
    #15 0x570d340 in llvm::legacy::PassManagerImpl::run(llvm::Module&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/LegacyPassManager.cpp:542:44
    #16 0x378bab6 in compileModule(char**, llvm::LLVMContext&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/llc/llc.cpp:687:8
    #17 0x37899ab in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/llc/llc.cpp:388:22
    #18 0x7f57aaa7309a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a)
    #19 0x37629b9 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/llc+0x37629b9)
SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp:12208:58 in 
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/CodeGen/AArch64/urem-seteq-nonzero.ll
--
********************
Testing:  0.. 10.. 20.. 30.. 40.. 50.. 60.. 70.. 80.. 90.. 
********************
Failed Tests (1):
  LLVM :: CodeGen/AArch64/urem-seteq-nonzero.ll

[benshi001]

@kda Could you please have a look at my new patch revision ? I have tested on my laptop (without sanitizer), which is slow to run full sanitizer tests.

[benshi001]

The overflow is caused by int64_t multiplication overflow, which I have changed to class APInt.

[kda]

In D108871#2983395, @benshi001 wrote:

@kda Could you please have a look at my new patch revision ? I have tested on my laptop (without sanitizer), which is slow to run full sanitizer tests.

Yes.
Appears to build correctly.
I apologize for the delay.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

bool lowerInterleavedLoad(LoadInst *LI,
                             ArrayRef<unsigned> Indices,
                             unsigned Factor) const override;
   bool lowerInterleavedStore(StoreInst *SI, ShuffleVectorInst *SVI,
                              unsigned Factor) const override;
 
   bool isLegalAddImmediate(int64_t) const override;
   bool isLegalICmpImmediate(int64_t) const override;
 
+  bool isMulAddWithConstProfitable(const SDValue &AddNode,
+                                   const SDValue &ConstNode) const override;
+
   bool shouldConsiderGEPOffsetSplit() const override;
 
   EVT getOptimalMemOpType(const MemOp &Op,
                           const AttributeList &FuncAttributes) const override;
 
   LLT getOptimalMemOpLLT(const MemOp &Op,
                          const AttributeList &FuncAttributes) const override;
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::isLegalAddImmediate(int64_t Immed) const {
   Immed = std::abs(Immed);
   bool IsLegal = ((Immed >> 12) == 0 ||
                   ((Immed & 0xfff) == 0 && Immed >> 24 == 0));
   LLVM_DEBUG(dbgs() << "Is " << Immed
                     << " legal add imm: " << (IsLegal ? "yes" : "no") << "\n");
   return IsLegal;
 }
 
+// Return true if Imm can be composed by a single MOVZ instruction.
+static inline bool isMOVZImm(uint64_t Imm, unsigned Sz) {
+  assert((Sz == 64) || (Sz == 32) && "invalid immediate size");
+  if (Sz == 32)
+    Imm &= 0xffffffff;
+  for (unsigned I = 0; I < Sz; I += 16) {
+    if ((Imm & ~0xffff) != 0)
+      return false;
+    Imm >>= 16;
+  }
+  return true;
+}
+
+// Return true if Imm can be composed by a single MOVN intruction.
+static inline bool isMOVNImm(uint64_t Imm, unsigned Sz) {
+  return isMOVZImm(~Imm, Sz);
+}
+
+// Return false to prevent folding
+// (mul (add x, c1), c2) -> (add (mul x, c2), c2*c1) in DAGCombine,
+// if the folding leads to worse code.
+bool AArch64TargetLowering::isMulAddWithConstProfitable(
+    const SDValue &AddNode, const SDValue &ConstNode) const {
+  // Let the DAGCombiner decide for vector types and large types.
+  const EVT VT = AddNode.getValueType();
+  if (VT.isVector() || VT.getScalarSizeInBits() > 64)
+    return true;
+
+  // It is worse if c1 is legal add immediate, while c1*c2 is not,
+  // and has to be composed by at least two instructions. (That means c1*c2
+  // is neither legal logical immediate nor composed by a single MOVZ/MOVN.)
+  const ConstantSDNode *C1Node = cast<ConstantSDNode>(AddNode.getOperand(1));
+  const ConstantSDNode *C2Node = cast<ConstantSDNode>(ConstNode);
+  const int64_t C1 = C1Node->getSExtValue();
+  const int64_t C2 = C2Node->getSExtValue();
+  const unsigned Sz = VT.getScalarSizeInBits() > 32 ? 64 : 32;
+  if (isLegalAddImmediate(C1) && !isLegalAddImmediate(C1 * C2) &&
+      !isMOVZImm(C1 * C2, Sz) && !isMOVNImm(C1 * C2, Sz) &&
+      !AArch64_AM::isLogicalImmediate(C1 * C2, Sz))
+    return false;
+
+  // Default to true and let the DAGCombiner decide.
+  return true;
+}
+
 // Integer comparisons are implemented with ADDS/SUBS, so the range of valid
 // immediates is the same as for an add or a sub.
 bool AArch64TargetLowering::isLegalICmpImmediate(int64_t Immed) const {
   return isLegalAddImmediate(Immed);
 }
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.

llvm/test/CodeGen/AArch64/addimm-mulimm.ll

; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, -31000
   %tmp1 = mul i32 %tmp0, 37
   ret i32 %tmp1
 }
 
 define i64 @addimm_mulimm_reject_00(i64 %a) {
 ; CHECK-LABEL: addimm_mulimm_reject_00:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #1200
-; CHECK-NEXT:    mov w8, #3700
-; CHECK-NEXT:    movk w9, #175, lsl #16
-; CHECK-NEXT:    madd x0, x0, x8, x9
+; CHECK-NEXT:    add x8, x0, #3100
+; CHECK-NEXT:    mov w9, #3700
+; CHECK-NEXT:    mul x0, x8, x9
 ; CHECK-NEXT:    ret
   %tmp0 = add i64 %a, 3100
   %tmp1 = mul i64 %tmp0, 3700
   ret i64 %tmp1
 }
 
 define i64 @addimm_mulimm_reject_01(i64 %a) {
 ; CHECK-LABEL: addimm_mulimm_reject_01:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x9, #-1200
-; CHECK-NEXT:    mov w8, #3700
-; CHECK-NEXT:    movk x9, #65360, lsl #16
-; CHECK-NEXT:    madd x0, x0, x8, x9
+; CHECK-NEXT:    sub x8, x0, #3100
+; CHECK-NEXT:    mov w9, #3700
+; CHECK-NEXT:    mul x0, x8, x9
 ; CHECK-NEXT:    ret
   %tmp0 = add i64 %a, -3100
   %tmp1 = mul i64 %tmp0, 3700
   ret i64 %tmp1
 }
 
 define signext i32 @addimm_mulimm_reject_02(i32 signext %a) {
 ; CHECK-LABEL: addimm_mulimm_reject_02:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #1200
-; CHECK-NEXT:    mov w8, #3700
-; CHECK-NEXT:    movk w9, #175, lsl #16
-; CHECK-NEXT:    madd w0, w0, w8, w9
+; CHECK-NEXT:    add w8, w0, #3100
+; CHECK-NEXT:    mov w9, #3700
+; CHECK-NEXT:    mul w0, w8, w9
 ; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, 3100
   %tmp1 = mul i32 %tmp0, 3700
   ret i32 %tmp1
 }
 
 define signext i32 @addimm_mulimm_reject_03(i32 signext %a) {
 ; CHECK-LABEL: addimm_mulimm_reject_03:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #64336
-; CHECK-NEXT:    mov w8, #3700
-; CHECK-NEXT:    movk w9, #65360, lsl #16
-; CHECK-NEXT:    madd w0, w0, w8, w9
+; CHECK-NEXT:    sub w8, w0, #3100
+; CHECK-NEXT:    mov w9, #3700
+; CHECK-NEXT:    mul w0, w8, w9
 ; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, -3100
   %tmp1 = mul i32 %tmp0, 3700
   ret i32 %tmp1
 }

llvm/test/CodeGen/AArch64/urem-seteq-nonzero.ll

; CHECK-NEXT:    ret
   %urem = urem i32 %X, 6
   %cmp = icmp eq i32 %urem, 3
   ret i1 %cmp
 }
 
 define i1 @t32_6_4(i32 %X) nounwind {
 ; CHECK-LABEL: t32_6_4:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #43691
-; CHECK-NEXT:    mov w9, #21844
-; CHECK-NEXT:    movk w8, #43690, lsl #16
-; CHECK-NEXT:    movk w9, #21845, lsl #16
-; CHECK-NEXT:    madd w8, w0, w8, w9
+; CHECK-NEXT:    mov w9, #43691
+; CHECK-NEXT:    sub w8, w0, #4
+; CHECK-NEXT:    movk w9, #43690, lsl #16
+; CHECK-NEXT:    mul w8, w8, w9
 ; CHECK-NEXT:    mov w9, #43690
 ; CHECK-NEXT:    ror w8, w8, #1
 ; CHECK-NEXT:    movk w9, #10922, lsl #16
 ; CHECK-NEXT:    cmp w8, w9
 ; CHECK-NEXT:    cset w0, lo
 ; CHECK-NEXT:    ret
   %urem = urem i32 %X, 6
   %cmp = icmp eq i32 %urem, 4
   ret i1 %cmp
 }
 
 define i1 @t32_6_5(i32 %X) nounwind {
 ; CHECK-LABEL: t32_6_5:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #43691
-; CHECK-NEXT:    mov w9, #43689
-; CHECK-NEXT:    movk w8, #43690, lsl #16
+; CHECK-NEXT:    mov w9, #43691
+; CHECK-NEXT:    sub w8, w0, #5
 ; CHECK-NEXT:    movk w9, #43690, lsl #16
-; CHECK-NEXT:    madd w8, w0, w8, w9
+; CHECK-NEXT:    mul w8, w8, w9
 ; CHECK-NEXT:    mov w9, #43690
 ; CHECK-NEXT:    ror w8, w8, #1
 ; CHECK-NEXT:    movk w9, #10922, lsl #16
 ; CHECK-NEXT:    cmp w8, w9
 ; CHECK-NEXT:    cset w0, lo
 ; CHECK-NEXT:    ret
   %urem = urem i32 %X, 6
   %cmp = icmp eq i32 %urem, 5