[AArch64][GlobalISel] Relax oneuse restriction for PTR_ADD chain combining to check addressing legality.

[aemerson]

Summary

Fixing this after @foad reported that the oneuse restriction regressed some code. I couldn't reproduce this without essentially disabling reassociation since that ends up optimizing the test case anyway.

[aemerson]

@foad does this fix your internal regression?

[foad]

In D105676#2866242, @aemerson wrote:

@foad does this fix your internal regression?

@sebastian-ne can you try it out please?

[sebastian-ne]

does this fix your internal regression?

Yes, thanks a lot.
FWIW, the “internal regression” is D103326, all the test changes there were gone after I rebased it upon D105069.

There is now only a single case, where slightly worse code than is generated compared to before D105069: https://github.com/llvm/llvm-project/blob/c282d55a38577e076b48cd7a8113e5eb0a2039cd/llvm/test/CodeGen/AMDGPU/GlobalISel/extractelement.ll#L4151-L4164
isLegalAddressingMode correctly returns false and the dead loads from the <64 x i32> load (which are multiple users of the G_PTR_ADD) are only removed later, after the G_PTR_ADD has long been converted to pseudo instructions. I guess there’s not much we can do about that.

[sebastian-ne]

Thinking more about it, maybe it makes sense to only skip the combine if isLegalAddressingMode(MaybeImm2Val->Value) && !isLegalAddressingMode(CombinedImm)?
(That should also fix the case I mentioned in the last comment.)

[sebastian-ne]

gentle reverse ping

[sebastian-ne]

Add check that addressing mode was legal before, but is not legal anymore after adding the offsets.

[aemerson]

llvm/test/CodeGen/AArch64/GlobalISel/combine-ptradd-reassociation.mir

26–27

I think this test is actually wrong, not sure how it slipped through. To properly test the reassociation behavior we should use a legal offset. Can you replace the 4777 with 1600 and the 6 with 3?

176–177

Same here.

[aemerson]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1731–1734

Also, we now have the GLoadStore wrapper class that can be used here with a dyn_cast.

1735

and this can then be replaced with LdStMI.getReg(0)

[aemerson]

Hope you don't mind me taking back this revision to address my comments and commit it.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

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;
 
-  // Don't do this combine if there multiple uses of the first PTR_ADD,
-  // since we may be able to compute the second PTR_ADD as an immediate
-  // offset anyway. Folding the first offset into the second may cause us
-  // to go beyond the bounds of our legal addressing modes.
-  if (!MRI.hasOneNonDBGUse(Add2))
-    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;
 
+  // Check if the new combined immediate forms an illegal addressing mode.
+  // Do not combine if it was legal before but would get illegal.
+  // To do so, we need to find a load/store user of the pointer to get
+  // the access type.
+  Type *AccessTy = nullptr;
+  auto &MF = *MI.getMF();
+  for (auto &UseMI : MRI.use_nodbg_instructions(MI.getOperand(0).getReg())) {
+    if (auto *LdSt = dyn_cast<GLoadStore>(&UseMI)) {
+      AccessTy = getTypeForLLT(MRI.getType(LdSt->getReg(0)),
+                               MF.getFunction().getContext());
+      break;
+    }
+  }
+  TargetLoweringBase::AddrMode AMNew;
+  APInt CombinedImm = MaybeImmVal->Value + MaybeImm2Val->Value;
+  AMNew.BaseOffs = CombinedImm.getSExtValue();
+  if (AccessTy) {
+    AMNew.HasBaseReg = true;
+    TargetLoweringBase::AddrMode AMOld;
+    AMOld.BaseOffs = MaybeImm2Val->Value.getSExtValue();
+    AMOld.HasBaseReg = true;
+    unsigned AS = MRI.getType(Add2).getAddressSpace();
+    const auto &TLI = *MF.getSubtarget().getTargetLowering();
+    if (TLI.isLegalAddressingMode(MF.getDataLayout(), AMOld, AccessTy, AS) &&
+        !TLI.isLegalAddressingMode(MF.getDataLayout(), AMNew, AccessTy, AS))
+      return false;
+  }
+
   // Pass the combined immediate to the apply function.
-  MatchInfo.Imm = (MaybeImmVal->Value + MaybeImm2Val->Value).getSExtValue();
+  MatchInfo.Imm = AMNew.BaseOffs;
   MatchInfo.Base = Base;
   return true;
 }
 
 void CombinerHelper::applyPtrAddImmedChain(MachineInstr &MI,
                                            PtrAddChain &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_PTR_ADD && "Expected G_PTR_ADD");
   MachineIRBuilder MIB(MI);

llvm/test/CodeGen/AArch64/GlobalISel/combine-ptradd-reassociation.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
 ---
 name:            test1_noreassoc_legal_already_new_is_illegal
 alignment:       4
 tracksRegLiveness: true
 liveins:
   - { reg: '$x0' }
 body:             |
   bb.1:
     liveins: $x0
 
     ; CHECK-LABEL: name: test1_noreassoc_legal_already_new_is_illegal
     ; CHECK: liveins: $x0
     ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
-    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 4777
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 1600
     ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 6
     ; CHECK: [[C2:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
     ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
     ; CHECK: [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD [[PTR_ADD]], [[C1]](s64)
     ; CHECK: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[PTR_ADD1]](p0) :: (load (s32))
     ; CHECK: G_STORE [[C2]](s32), [[PTR_ADD]](p0) :: (store (s32))
     ; CHECK: $w0 = COPY [[LOAD]](s32)
     ; CHECK: RET_ReallyLR implicit $w0
     %0:_(p0) = COPY $x0
-    %2:_(s64) = G_CONSTANT i64 4777
+    %2:_(s64) = G_CONSTANT i64 1600
     %4:_(s64) = G_CONSTANT i64 6
     %9:_(s32) = G_CONSTANT i32 0
     %10:_(p0) = G_PTR_ADD %0, %2(s64)
     %11:_(p0) = G_PTR_ADD %10, %4(s64)
     %7:_(s32) = G_LOAD %11(p0) :: (load 4)
     G_STORE %9(s32), %10(p0) :: (store 4) ; other use of %10
     $w0 = COPY %7(s32)
     RET_ReallyLR implicit $w0
 

 alignment:       4
 tracksRegLiveness: true
 body:             |
   bb.1:
     liveins: $x0
     ; CHECK-LABEL: name: walk_through_inttoptr
     ; CHECK: liveins: $x0
     ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
-    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 4777
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 1600
     ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 6
     ; CHECK: [[C2:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
     ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
     ; CHECK: %ptr_to_int:_(s64) = G_PTRTOINT [[PTR_ADD]](p0)
     ; CHECK: [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD [[PTR_ADD]], [[C1]](s64)
     ; CHECK: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[PTR_ADD1]](p0) :: (load (s32))
     ; CHECK: G_STORE [[C2]](s32), [[PTR_ADD]](p0) :: (store (s32))
     ; CHECK: G_STORE %ptr_to_int(s64), [[PTR_ADD]](p0) :: (store (s64))
     ; CHECK: $w0 = COPY [[LOAD]](s32)
     ; CHECK: RET_ReallyLR implicit $w0
     %0:_(p0) = COPY $x0
-    %2:_(s64) = G_CONSTANT i64 4777
+    %2:_(s64) = G_CONSTANT i64 1600
     %4:_(s64) = G_CONSTANT i64 6
     %9:_(s32) = G_CONSTANT i32 0
     %10:_(p0) = G_PTR_ADD %0, %2(s64)
     %ptr_to_int:_(s64) = G_PTRTOINT %10
     %11:_(p0) = G_PTR_ADD %10, %4(s64)
     %7:_(s32) = G_LOAD %11(p0) :: (load 4)
     G_STORE %9(s32), %10(p0) :: (store 4) ; other use of %10
     G_STORE %ptr_to_int(s64), %10(p0) :: (store 8)
     $w0 = COPY %7(s32)
     RET_ReallyLR implicit $w0

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
+# RUN: llc -mtriple aarch64-apple-ios  -run-pass=aarch64-prelegalizer-combiner --aarch64prelegalizercombinerhelper-only-enable-rule="ptr_add_immed_chain"  %s -o - -verify-machineinstrs | FileCheck %s
+# REQUIRES: asserts
 
 # 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
bb.1:
     %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
 ...
+---
+name: ptradd_would_form_illegal_load_addressing
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: ptradd_would_form_illegal_load_addressing
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 4
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 4096
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
+    ; CHECK: [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD [[PTR_ADD]], [[C1]](s64)
+    ; CHECK: %ld:_(s64) = G_LOAD [[PTR_ADD1]](p0) :: (load (s64))
+    ; CHECK: %ld_other:_(s64) = G_LOAD [[PTR_ADD]](p0) :: (load (s64))
+    ; CHECK: $x0 = COPY %ld(s64)
+    ; CHECK: $x1 = COPY %ld_other(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_CONSTANT i64 4
+    %2:_(s64) = G_CONSTANT i64 4096
+    %3:_(p0) = G_PTR_ADD %0(p0), %1
+    %4:_(p0) = G_PTR_ADD %3(p0), %2
+    %ld:_(s64) = G_LOAD %4(p0) :: (load 8)
+    %ld_other:_(s64) = G_LOAD %3(p0) :: (load 8)
+    $x0 = COPY %ld(s64)
+    $x1 = COPY %ld_other(s64)
+    RET_ReallyLR implicit $x0
+...