Allow rematerialization of ARM Thumb MOVi8 instruction in some contexts

[philipginsbach]

Summary

Constants are crucial for code size in the ARM Thumb-1 instruction set.
The 16 bit instruction size often does not offer enough space for immediate arguments.
This means that additional instructions are frequently used to load constants into registers.
Since constants are hoisted, this can lead to significant register spillage if they are used multiple times in a single function.
This can be avoided by rematerialization, i.e. recomputing a constant instead of reloading it from the stack.
This patch extends the rematerialization functionality of llvm such that it is able to rematerialize ARM Thumb MOVSi8 instruction despite it setting the CPSR flags.

[samparker]

Hi Philip,

This needs some tests.

cheers,
sam

[samparker]

lib/CodeGen/LiveRangeEdit.cpp

162

For speed, I think you can iterate over just the defs instead of all the operands, they'll be fewer and it will negate the need to do the Reg and Def checks.

177

Allocatable isn't about liveliness, it's just whether the register can be used by the allocator generally, I'm guessing that this works here because it can't allocate the CPSR.

181

Can you not just query IsDead on the operand instead?

[philipginsbach]

lib/CodeGen/LiveRangeEdit.cpp

162

Agreed, I will fix this.

177

Yes, I'm aware of that.
I discussed that part with James Molloy in quite some depth.
From what I understand (and as I say in the comment there might be gaps), I can't rely on liveliness computation here, because further iterations of the register allocator might change the liveliness.
Obviously this problem goes away if a register can't be used freely by the register allocator.

Long story short: I'm quite certain the condition I check here is sufficient but it is probably way too strong.

181

I tried several different options here.
If I recall correctly, what you propose only works for virtual registers.

[samparker]

lib/Target/ARM/ARMInstrThumb.td

1136

Can this also be used for the movs register variant?

[philipginsbach]

lib/CodeGen/LiveRangeEdit.cpp

177

The important thing is that the patch is targeted at CPSR, as you point out, which is not allocatable.

[samparker]

lib/CodeGen/LiveRangeEdit.cpp

177

Ok, definitely better to be conservative then.

[philipginsbach]

lib/Target/ARM/ARMInstrThumb.td

1136

The infrastructure that the patch sets up could be used for much more generic instructions, yes!
Potentially, add, sub etc could be rematerialized the same way as well.

Unfortunately it's pointless though.

The problem is that the used registers would have to be preserved by chance from the first invocation to the second.
The register allocator is not incentivized to make that happen, so I it's extremely unlikely to occur often, especially since the whole thing becomes relevant only in a high register pressure situation.

[qcolombet]

Hi Philip,

Thanks for looking into extending the rematerialization support.

I believe the interference checking is too weak to actually work in all cases (see inline comments). More over, I think we cannot say this is still trivially rematerializable given now the callers of that method need to do interference checks. Thus I would suggest to use a new/different hook. In particular, I am not convinced the register coalescer will do the right thing if it rematerializes an instruction with more than one def.

Cheers,
-Quentin

include/llvm/Target/TargetInstrInfo.h

111

No else after return per LLVM coding standard.

lib/CodeGen/LiveRangeEdit.cpp

162

This is wrong. This misses all the implicit-def.

>

/// Returns a range over all explicit operands that are register definitions.
/// Implicit definition are not included!

166

That doesn't seem right either. Indeed where is the check for virtregs?

E.g., consider
vreg0,vreg1 = loadimm
...
vreg1 = something else
/// <=== if you remat here you're going to break the next use of vreg1

vreg0

vreg1

test/CodeGen/Thumb/movi8remat.ll

1 ↗

(On Diff #106606)

Please use a .mir test.
You can generate this with -stop-before greedy -simplify-mir

[philipginsbach]

Hi Quentin,

Thanks a lot for the comments, that's really valuable feedback.
I will take some time to dig through those issues and get back to you.

Best,
Philip

[philipginsbach]

Hi Quentin,

Let me follow up on your comments.

I renamed the member functions to isPotentiallyTriviallyReMaterializable and isReallyPotentiallyTriviallyReMaterializableGeneric.
This is somewhat grotesque but does bring across the functionality more clearly I think. Obviously I am open to more aesthetic suggestions.

The problem with implicit defs should be removed now and was not present in a previous version of the patch.
I didn't pay attention and was mislead by misleading documentation for MachineInstr::getNumOperands.

I think your concern for checking virtual register defs stems from isPotentiallyTriviallyReMaterializable not being very clear in it's purpose.
I hope the change to the return enum value improves this. Do you agree that this is sufficient or am I overlooking something crucial here?

Best,
Philip

include/llvm/Target/TargetInstrInfo.h

111

I will fix this.

lib/CodeGen/LiveRangeEdit.cpp

162

I was too eager to change this, I think the previous version should have been correct.

There appears to be a mistake in the documentation for MachineInstr::getNumOperands, it says "Access to explicit operands of the instruction."
However it includes the implicit operands.

166

I did not weaken the check for additional virtual registers in TargetInstrInfo::isReallyTriviallyReMaterializableGeneric.
What you describe would still be filtered out by that function (i.e. return Rematerializability::NO).
The only way in that it is more lenient than before is that is allows additional explicit physical register defs.

Arguably the behavior that you imply (allowing additional virtreg defs) would be useful and more obvious but at the core this patch is concerned with the movi8 instruction and the current change is enough to rematerialize that.

I will change the enum item YES_BUT_EXTRA_DEFS to YES_BUT_EXTRA_PHYSREG_DEFS to make the behavior of the function more obvious.

test/CodeGen/Thumb/movi8remat.ll

1 ↗

(On Diff #106606)

Ok, I will change it to a .mir test, thanks for providing the options.

[philipginsbach]

I have rebased onto the master branch at 244c796c894f50fb53f9dbe7627702661dfe69c2.

Quentin, you gave many useful suggestions before, what do you think is the way forward with this?

[dmgreen]

I ran some benchmarks for thumb1, they look great.

[zzheng]

lib/CodeGen/LiveRangeEdit.cpp

165

shouldn't we use auto type and range-based loop now?

[philipginsbach]

Uses range based loop now.

[philipginsbach]

lib/CodeGen/LiveRangeEdit.cpp

165

Thanks, I fixed it.

[philipginsbach]

I think I have addressed all your previous concerns, @qcolombet, is there anything else you'd like to see clarified?
Sam Parker suggested I add you as a reviewer, @MatzeB, do you have any suggestions?

Sorry to be pushy, this patch has been lying around for a while and it would be quite satisfying to get it through the process.

[qcolombet]

Hi,

Sorry for the delay, I haven’t had time to review the patch but I have one high level comment.
I am not sure the enum approach is desirable as it basically hard code the kind of constraints we can report and will grow very quickly if we want to extend it.
For instance, let say that on top of physreg defs we want to report virtual reg defs, now we would need “yes”, “no”, “yes_phys”, “yes_virt”, “yes_phys_virt” and the list grows with the cross product of everything we may want/need to report.

Anyhow, I’ll look closer at the patch to give more concrete feedbacks.

Cheers,
Q

[philipginsbach]

I am not sure the enum approach is desirable as it basically hard code the kind of constraints we can report and will grow very quickly if we want to extend it.
For instance, let say that on top of physreg defs we want to report virtual reg defs, now we would need “yes”, “no”, “yes_phys”, “yes_virt”, “yes_phys_virt” and the list grows with the cross product of everything we may want/need to report.

Hi Quentin,

Thanks a lot for looking at this again.

I understand your concern, the easiest way of tackling it would probably be some kind of bitfield with flags for the different kind of (potential) extensions of the unrestricted rematerializability property.

However, I would vote for that to be introduced when the problem actually arises. It isn't so obvious to me that many more generalizations of rematerializability will be necessary and the path from enum to bitfield is easy.

What do you think?

Cheers,
Philip

[qcolombet]

However, I would vote for that to be introduced when the problem actually arises. It isn't so obvious to me that many more generalizations of rematerializability will be necessary and the path from enum to bitfield is easy.

What do you think?

I agree, moving with incremental steps is the right thing to do.
The question is more does this step go in the direction we want? (I am not saying it doesn't, I just haven't looked closer enough to have an opinion :))

Now, I just wanted to point out that there are more opportunities to generalize the rematerialization. The obvious one to me is rematerializing everything (like pulling chain of computations, full instructions (with both definitions and arguments)) and in that sense, only the instruction or chain of instructions carry the right level of information. Right now we often introduce pseudo instruction to work around this limitation and that's what I would like to solve.

Anyway, I'll take a closer look at the patch.

[philipginsbach]

Now, I just wanted to point out that there are more opportunities to generalize the rematerialization. The obvious one to me is rematerializing everything (like pulling chain of computations, full instructions (with both definitions and arguments)) and in that sense, only the instruction or chain of instructions carry the right level of information. Right now we often introduce pseudo instruction to work around this limitation and that's what I would like to solve.

It's true, rematerializing a whole chain of instructions would require a more thorough approach.
I think it's important to be clear about motivations beforehand though: this patch was done to optimize code size for ARM Thumb. Since spilling and reloading takes only two instructions, rematerializing even chains of two instructions is pointless from that perspective. Obviously, there are other metrics to consider, but I think it would be important nonetheless to establish what exactly is to gain.

include/llvm/Target/TargetInstrInfo.h

                                          const TargetRegisterInfo *TRI,
                                          const MachineFunction &MF) const;
 
+  /// An instruction that pollutes additional registers might still be
+  /// rematerializable under the assumption that those registers aren't live.
+  /// This is the purpose of YES_BUT_EXTRA_DEFS.
+  enum class Rematerializability
+  {
+      NO = 0,
+      YES_BUT_EXTRA_DEFS = 1,
+      YES = 2
+  };
+
   /// Return true if the instruction is trivially rematerializable, meaning it
   /// has no side effects and requires no operands that aren't always available.
   /// This means the only allowed uses are constants and unallocatable physical
   /// registers so that the instructions result is independent of the place
   /// in the function.
-  bool isTriviallyReMaterializable(const MachineInstr &MI,
-                                   AliasAnalysis *AA = nullptr) const {
-    return MI.getOpcode() == TargetOpcode::IMPLICIT_DEF ||
-           (MI.getDesc().isRematerializable() &&
-            (isReallyTriviallyReMaterializable(MI, AA) ||
-             isReallyTriviallyReMaterializableGeneric(MI, AA)));
+  Rematerializability
+  isTriviallyReMaterializable(const MachineInstr &MI,
+                              AliasAnalysis *AA = nullptr) const {
+    if(MI.getOpcode() == TargetOpcode::IMPLICIT_DEF)
+        return Rematerializability::YES;
+    else if(!MI.getDesc().isRematerializable())
+        return Rematerializability::NO;
+    else if(isReallyTriviallyReMaterializable(MI, AA))
+        return Rematerializability::YES;
+    else
+        return isReallyTriviallyReMaterializableGeneric(MI, AA);
   }
 
 protected:

   /// set and the target hook isReallyTriviallyReMaterializable returns false,
   /// this function does target-independent tests to determine if the
   /// instruction is really trivially rematerializable.
-  bool isReallyTriviallyReMaterializableGeneric(const MachineInstr &MI,
-                                                AliasAnalysis *AA) const;
+  Rematerializability
+  isReallyTriviallyReMaterializableGeneric(const MachineInstr &MI,
+                                           AliasAnalysis *AA) const;
 
 public:
   /// These methods return the opcode of the frame setup/destroy instructions

lib/CodeGen/CalcSpillWeights.cpp

       }
     }
 
-    if (!TII.isTriviallyReMaterializable(*MI, LIS.getAliasAnalysis()))
+    if (TII.isTriviallyReMaterializable(*MI, LIS.getAliasAnalysis())
+        == TargetInstrInfo::Rematerializability::NO)
       return false;
   }
   return true;

lib/CodeGen/LiveRangeEdit.cpp

                                           AliasAnalysis *aa) {
   assert(DefMI && "Missing instruction");
   ScannedRemattable = true;
-  if (!TII.isTriviallyReMaterializable(*DefMI, aa))
+  if (TII.isTriviallyReMaterializable(*DefMI, aa)
+      == TargetInstrInfo::Rematerializability::NO)
     return false;
   Remattable.insert(VNI);
   return true;

   if (cheapAsAMove && !TII.isAsCheapAsAMove(*RM.OrigMI))
     return false;
 
+  // The instruction passed the checkRematerializable criterions. Now
+  // that we know the context, we need to make sure the instruction does
+  // not def any additional live registers.
+  if (TII.isTriviallyReMaterializable(*RM.OrigMI, nullptr)
+      == TargetInstrInfo::Rematerializability::YES_BUT_EXTRA_DEFS)
+  {
+    for (unsigned i = 0, e = RM.OrigMI->getNumOperands(); i != e; ++i) {
+      const MachineOperand &MO = RM.OrigMI->getOperand(i);
+      if (!MO.isReg()) continue;
+      unsigned Reg = MO.getReg();
+
+      // Check for a well-behaved physical register.
+      if (TargetRegisterInfo::isPhysicalRegister(Reg) && MO.isDef()) {
+        // A physreg def. We need to make sure the register is dead.
+        SlotIndexes       *Indexes;
+        MachineInstr      *Instruction;
+        MachineBasicBlock *BasicBlock;
+
+        // The interaction with the register allocator isn't entirely clear
+        // to me, so to be on the safe side, never assume registers to be
+        // dead if they are allocatable.
+        if (MRI.isAllocatable(Reg) ||
+            !(Indexes     = LIS.getSlotIndexes()) ||
+            !(Instruction = Indexes->getInstructionFromIndex(UseIdx)) ||
+            !(BasicBlock  = Instruction->getParent()) ||
+            (BasicBlock->computeRegisterLiveness(
+                MRI.getTargetRegisterInfo(), Reg, Instruction)
+              != MachineBasicBlock::LivenessQueryResult::LQR_Dead)) {
+          return false;
+        }
+      }
+    }
+  }
+
   // Verify that all used registers are available with the same values.
   if (!allUsesAvailableAt(RM.OrigMI, DefIdx, UseIdx))
     return false;

     // the inst for remat of other siblings. The inst is saved in
     // LiveRangeEdit::DeadRemats and will be deleted after all the
     // allocations of the func are done.
-    if (isOrigDef && DeadRemats && TII.isTriviallyReMaterializable(*MI, AA)) {
+    if (isOrigDef && DeadRemats &&
+        TII.isTriviallyReMaterializable(*MI, AA)
+        == TargetInstrInfo::Rematerializability::YES) {
       LiveInterval &NewLI = createEmptyIntervalFrom(Dest);
       NewLI.removeEmptySubRanges();
       VNInfo *VNI = NewLI.getNextValue(Idx, LIS.getVNInfoAllocator());

lib/CodeGen/MachineLICM.cpp

 
   // Rematerializable instructions should always be hoisted since the register
   // allocator can just pull them down again when needed.
-  if (TII->isTriviallyReMaterializable(MI, AA))
+  if (TII->isTriviallyReMaterializable(MI, AA)
+      != TargetInstrInfo::Rematerializability::NO)
     return true;
 
   // FIXME: If there are long latency loop-invariant instructions inside the

 
   // High register pressure situation, only hoist if the instruction is going
   // to be remat'ed.
-  if (!TII->isTriviallyReMaterializable(MI, AA) &&
+  if (TII->isTriviallyReMaterializable(MI, AA)
+      == TargetInstrInfo::Rematerializability::NO &&
       !MI.isDereferenceableInvariantLoad(AA)) {
     DEBUG(dbgs() << "Can't remat / high reg-pressure: " << MI);
     return false;

lib/CodeGen/RegisterCoalescer.cpp

   }
   if (!TII->isAsCheapAsAMove(*DefMI))
     return false;
-  if (!TII->isTriviallyReMaterializable(*DefMI, AA))
+  if (TII->isTriviallyReMaterializable(*DefMI, AA)
+      != TargetInstrInfo::Rematerializability::YES)
     return false;
   if (!definesFullReg(*DefMI, SrcReg))
     return false;

lib/CodeGen/TargetInstrInfo.cpp

   return NewMI;
 }
 
-bool TargetInstrInfo::isReallyTriviallyReMaterializableGeneric(
+TargetInstrInfo::Rematerializability
+TargetInstrInfo::isReallyTriviallyReMaterializableGeneric(
     const MachineInstr &MI, AliasAnalysis *AA) const {
   const MachineFunction &MF = *MI.getParent()->getParent();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
 
   // Remat clients assume operand 0 is the defined register.
   if (!MI.getNumOperands() || !MI.getOperand(0).isReg())
-    return false;
+    return Rematerializability::NO;
   unsigned DefReg = MI.getOperand(0).getReg();
 
   // A sub-register definition can only be rematerialized if the instruction

   // moved safely.
   if (TargetRegisterInfo::isVirtualRegister(DefReg) &&
       MI.getOperand(0).getSubReg() && MI.readsVirtualRegister(DefReg))
-    return false;
+    return Rematerializability::NO;
 
   // A load from a fixed stack slot can be rematerialized. This may be
   // redundant with subsequent checks, but it's target-independent,

   int FrameIdx = 0;
   if (isLoadFromStackSlot(MI, FrameIdx) &&
       MF.getFrameInfo().isImmutableObjectIndex(FrameIdx))
-    return true;
+    return Rematerializability::YES;
 
   // Avoid instructions obviously unsafe for remat.
   if (MI.isNotDuplicable() || MI.mayStore() || MI.hasUnmodeledSideEffects())
-    return false;
+    return Rematerializability::NO;
 
   // Don't remat inline asm. We have no idea how expensive it is
   // even if it's side effect free.
   if (MI.isInlineAsm())
-    return false;
+    return Rematerializability::NO;
 
   // Avoid instructions which load from potentially varying memory.
   if (MI.mayLoad() && !MI.isDereferenceableInvariantLoad(AA))
-    return false;
+    return Rematerializability::NO;
+
+  // Track whether the instruction pollutes any additional registers.
+  // if they are dead at the rematerialization location, it's still ok.
+  bool AdditionalDefs = false;
 
   // If any of the registers accessed are non-constant, conservatively assume
   // the instruction is not rematerializable.

         // and we can freely move its uses. Alternatively, if it's allocatable,
         // it could get allocated to something with a def during allocation.
         if (!MRI.isConstantPhysReg(Reg))
-          return false;
+          return Rematerializability::NO;
       } else {
-        // A physreg def. We can't remat it.
-        return false;
+        // A physreg def. If the register is dead, we can still rematerialize.
+        // This will be checked in LiveRangeEdit::canRematerializeAt.
+        AdditionalDefs = true;
       }
       continue;
     }

     // Only allow one virtual-register def.  There may be multiple defs of the
     // same virtual register, though.
     if (MO.isDef() && Reg != DefReg)
-      return false;
+      return Rematerializability::NO;
 
     // Don't allow any virtual-register uses. Rematting an instruction with
     // virtual register uses would length the live ranges of the uses, which
     // is not necessarily a good idea, certainly not "trivial".
     if (MO.isUse())
-      return false;
+      return Rematerializability::NO;
   }
 
   // Everything checked out.
-  return true;
+  return AdditionalDefs ? Rematerializability::YES_BUT_EXTRA_DEFS
+                        : Rematerializability::YES;
 }
 
 int TargetInstrInfo::getSPAdjust(const MachineInstr &MI) const {

lib/Target/ARM/ARMInstrThumb.td

                 [(set tGPR:$Rdn, (srl tGPR:$Rn, tGPR:$Rm))]>, Sched<[WriteALU]>;
 
 // Move register
-let isMoveImm = 1 in
+let isMoveImm = 1, isReMaterializable = 1 in
 def tMOVi8 : T1sI<(outs tGPR:$Rd), (ins imm0_255:$imm8), IIC_iMOVi,
                   "mov", "\t$Rd, $imm8",
                   [(set tGPR:$Rd, imm0_255:$imm8)]>,