[AArch64][GlobalISel] Fold a chain of two G_PTR_ADDs of constant offsets.

[aemerson]

Summary

E.g.

%addr1 = G_PTR_ADD %base, G_CONSTANT 20
%addr2 = G_PTR_ADD %addr1, G_CONSTANT 8
    -->
%addr2 = G_PTR_ADD %base, G_CONSTANT 28

[paquette]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1356

Why not getConstantVRegValWithLookThrough?

1366

Why not getConstantVRegValWithLookThrough?

[aemerson]

Use getConstantVRegValWithLookThrough.

[paquette]

LGTM

[arsenm]

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-ptradd-chain.mir

27

Could use a case that uses the look through, and one with a second pair to fold

[aemerson]

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-ptradd-chain.mir

27

Ok, will add those tests when I commit.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

 
 struct IndexedLoadStoreMatchInfo {
   Register Addr;
   Register Base;
   Register Offset;
   bool IsPre;
 };
 
+struct PtrAddChain {
+  int64_t Imm;
+  Register Base;
+};
+
 class CombinerHelper {
 protected:
   MachineIRBuilder &Builder;
   MachineRegisterInfo &MRI;
   GISelChangeObserver &Observer;
   GISelKnownBits *KB;
   MachineDominatorTree *MDT;
 
public:
   ///
   /// -->
   ///
   ///     $addr = G_INDEXED_STORE $val, $base, $offset
   ///     [...]
   ///     $whatever = COPY $addr
   bool tryCombineMemCpyFamily(MachineInstr &MI, unsigned MaxLen = 0);
 
+  bool matchPtrAddImmedChain(MachineInstr &MI, PtrAddChain &MatchInfo);
+  bool applyPtrAddImmedChain(MachineInstr &MI, PtrAddChain &MatchInfo);
+
   /// Try to transform \p MI by using all of the above
   /// combine functions. Returns true if changed.
   bool tryCombine(MachineInstr &MI);
 
 private:
   // Memcpy family optimization helpers.
   bool optimizeMemcpy(MachineInstr &MI, Register Dst, Register Src,
                       unsigned KnownLen, unsigned DstAlign, unsigned SrcAlign,

llvm/include/llvm/Target/GlobalISel/Combine.td

 // FIXME: Is there a reason this wasn't in tryCombine? I've left it out of
 //        all_combines because it wasn't there.
 def elide_br_by_inverting_cond : GICombineRule<
   (defs root:$root),
   (match (wip_match_opcode G_BR):$root,
          [{ return Helper.matchElideBrByInvertingCond(*${root}); }]),
   (apply [{ Helper.applyElideBrByInvertingCond(*${root}); }])>;
 
-def all_combines : GICombineGroup<[trivial_combines, combines_for_extload, combine_indexed_load_store]>;
+def ptr_add_immed_matchdata : GIDefMatchData<"PtrAddChain">;
+def ptr_add_immed_chain : GICombineRule<
+  (defs root:$d, ptr_add_immed_matchdata:$matchinfo),
+  (match (wip_match_opcode G_PTR_ADD):$d,
+         [{ return Helper.matchPtrAddImmedChain(*${d}, ${matchinfo}); }]),
+  (apply [{ Helper.applyPtrAddImmedChain(*${d}, ${matchinfo}); }])>;
+
+def all_combines : GICombineGroup<[trivial_combines, ptr_add_immed_chain,
+    combines_for_extload, combine_indexed_load_store]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

if (ID == Intrinsic::memcpy)
     return optimizeMemcpy(MI, Dst, Src, KnownLen, DstAlign, SrcAlign, IsVolatile);
   if (ID == Intrinsic::memmove)
     return optimizeMemmove(MI, Dst, Src, KnownLen, DstAlign, SrcAlign, IsVolatile);
   if (ID == Intrinsic::memset)
     return optimizeMemset(MI, Dst, Src, KnownLen, DstAlign, IsVolatile);
   return false;
 }
 
+bool CombinerHelper::matchPtrAddImmedChain(MachineInstr &MI,
+                                           PtrAddChain &MatchInfo) {
+  // We're trying to match the following pattern:
+  //   %t1 = G_PTR_ADD %base, G_CONSTANT imm1
+  //   %root = G_PTR_ADD %t1, G_CONSTANT imm2
+  // -->
+  //   %root = G_PTR_ADD %base, G_CONSTANT (imm1 + imm2)
+
+  if (MI.getOpcode() != TargetOpcode::G_PTR_ADD)
+    return false;
+
+  Register Add2 = MI.getOperand(1).getReg();
+  Register Imm1 = MI.getOperand(2).getReg();
+  auto MaybeImmVal = getConstantVRegValWithLookThrough(Imm1, MRI);
+  if (!MaybeImmVal)
+    return false;
+
+  MachineInstr *Add2Def = MRI.getUniqueVRegDef(Add2);
+  if (!Add2Def || Add2Def->getOpcode() != TargetOpcode::G_PTR_ADD)
+    return false;
+
+  Register Base = Add2Def->getOperand(1).getReg();
+  Register Imm2 = Add2Def->getOperand(2).getReg();
+  auto MaybeImm2Val = getConstantVRegValWithLookThrough(Imm2, MRI);
+  if (!MaybeImm2Val)
+    return false;
+
+  // Pass the combined immediate to the apply function.
+  MatchInfo.Imm = MaybeImmVal->Value + MaybeImm2Val->Value;
+  MatchInfo.Base = Base;
+  return true;
+}
+
+bool CombinerHelper::applyPtrAddImmedChain(MachineInstr &MI,
+                                           PtrAddChain &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_PTR_ADD && "Expected G_PTR_ADD");
+  MachineIRBuilder MIB(MI);
+  LLT OffsetTy = MRI.getType(MI.getOperand(2).getReg());
+  auto NewOffset = MIB.buildConstant(OffsetTy, MatchInfo.Imm);
+  Observer.changingInstr(MI);
+  MI.getOperand(1).setReg(MatchInfo.Base);
+  MI.getOperand(2).setReg(NewOffset.getReg(0));
+  Observer.changedInstr(MI);
+  return true;
+}
+
 bool CombinerHelper::tryCombine(MachineInstr &MI) {
   if (tryCombineCopy(MI))
     return true;
   if (tryCombineExtendingLoads(MI))
     return true;
   if (tryCombineIndexedLoadStore(MI))
     return true;
   return false;
 }

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-ptradd-chain.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64-apple-ios  -run-pass=aarch64-prelegalizer-combiner %s -o - -verify-machineinstrs | FileCheck %s
+
+# Check that we fold two adds of constant offsets with G_PTR_ADD into a single G_PTR_ADD.
+---
+name: ptradd_chain
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: ptradd_chain
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
+    ; CHECK: $x0 = COPY [[PTR_ADD]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_CONSTANT i64 4
+    %2:_(s64) = G_CONSTANT i64 12
+    %3:_(p0) = G_PTR_ADD %0(p0), %1
+    %4:_(p0) = G_PTR_ADD %3(p0), %2
+    $x0 = COPY %4(p0)
+    RET_ReallyLR implicit $x0
+...
+---
+name: ptradd_chain_2
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0
+    ; CHECK-LABEL: name: ptradd_chain_2
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 28
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
+    ; CHECK: $x0 = COPY [[PTR_ADD]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_CONSTANT i64 4
+    %2:_(s64) = G_CONSTANT i64 12
+    %3:_(p0) = G_PTR_ADD %0(p0), %1
+    %4:_(p0) = G_PTR_ADD %3(p0), %2
+    %5:_(p0) = G_PTR_ADD %4(p0), %2
+    $x0 = COPY %5(p0)
+    RET_ReallyLR implicit $x0
+...
+---
+name: ptradd_chain_lookthough
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0
+    ; CHECK-LABEL: name: ptradd_chain_lookthough
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 28
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
+    ; CHECK: $x0 = COPY [[PTR_ADD]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_CONSTANT i64 4
+    %2:_(s64) = G_CONSTANT i64 12
+    %6:_(s32) = G_TRUNC %2(s64)
+    %7:_(s64) = G_SEXT %6(s32)
+    %3:_(p0) = G_PTR_ADD %0(p0), %1
+    %4:_(p0) = G_PTR_ADD %3(p0), %2
+    %5:_(p0) = G_PTR_ADD %4(p0), %7
+    $x0 = COPY %5(p0)
+    RET_ReallyLR implicit $x0
+...