[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 ↗

(On Diff #362739)

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 ↗

(On Diff #362739)

Same here.

[aemerson]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1738–1741

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

1742

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.
+  // 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())) {
+    switch (UseMI.getOpcode()) {
+    case TargetOpcode::G_LOAD:
+    case TargetOpcode::G_ZEXTLOAD:
+    case TargetOpcode::G_SEXTLOAD:
+    case TargetOpcode::G_STORE:
+      AccessTy = getTypeForLLT(MRI.getType(UseMI.getOperand(0).getReg()),
+                               MF.getFunction().getContext());
+      break;
+    }
+  }
+  TargetLoweringBase::AddrMode AM;
+  APInt CombinedImm = MaybeImmVal->Value + MaybeImm2Val->Value;
+  AM.BaseOffs = CombinedImm.getSExtValue();
+  if (AccessTy) {
+    AM.HasBaseReg = true;
+    unsigned AS = MRI.getType(Add2).getAddressSpace();
+    const auto &TLI = *MF.getSubtarget().getTargetLowering();
+    if (!TLI.isLegalAddressingMode(MF.getDataLayout(), AM, AccessTy, AS))
+      return false;
+  }
+
   // Pass the combined immediate to the apply function.
-  MatchInfo.Imm = (MaybeImmVal->Value + MaybeImm2Val->Value).getSExtValue();
+  MatchInfo.Imm = AM.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/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
+...