[AIEX] Basic LICM pass for control registers #133

gbossu · 2024-07-22T09:55:42Z

This is reviving some older work I made to hoist assignments to CR registers out of loops. This can both hoist to the pre-header, or sink to the exit block. This is needed to clean up our loops so they can later be SW pipelined. QoR shows minor improvements, and one regression: Add2D. What happens is that we now SW-pipeline the inner loop, but spill vector registers. In a future PR, I'll address those spills by teaching the pre-RA scheduler to try even harder to limit register pressure in cases where it is especially high.

E.g.

loop:
  $cr = MOV ...
  USE $cr
  $cr = MOV 0
  JNZ ... loop

into:

  $cr = MOV ...
loop:
  USE $cr
  JNZ ... loop
exit:
  $cr = MOV 0

Limitations:

Only supports simple one-bb loops.

andcarminati · 2024-07-26T11:24:58Z

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

+
+using namespace llvm;
+
+#define DEBUG_TYPE "reserved-reg-licm"


nit: aie-reserved-reg-licm because this is also the name of the pass.

andcarminati · 2024-07-26T11:33:46Z

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

+static MCRegister getSinglePhysRegDef(const MachineInstr &MI) {
+  if (MI.getNumDefs() == 1 && MI.getNumExplicitDefs() == 1) {
+    const MachineOperand &MO = MI.getOperand(0);
+    return MO.isReg() && MO.getReg().isPhysical() ? MO.getReg().asMCReg()


Is there any possibility that MO is not a register? I think your first If guarantees this condition.

I think it that it can be simpler:

if (MO.getReg().isPhysical()) return MO.getReg().asMCReg();

andcarminati · 2024-07-26T11:55:18Z

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

+  }
+
+  // Traverse instructions to remove defs
+  for (const MachineInstr &MI : reverse(Header))


~~I was thinking about one case that one register is defined by a block that dominates the Header. I think this case is not captured here, this register will not be in the live in set.~~

~~Can we just trust LoopBody->liveins()? set~~

Unserstood now: Note we do not need to extend the liveins of the loops BBs with the // newly-hoisted def, because reserved regs are always considered live.

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

martien-de-jong · 2024-07-30T11:52:52Z

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

+    if (MRI->canSimplifyPhysReg(PhysReg)) {
+      LiveRegs.addReg(PhysReg);
+    }
+  }


I would make this function take a block, not a loop, to make it easier to reuse in multi-block loops. And perhaps split off livereg initialization, which would be different for exiting blocks and internal blocks.

I remember in my initial work I extended the pass to support multi-BB loops. Typically outer loops. It was enough to only look in the header, and be conservative if a reserved register was not met there. If I need something more complicated, I'd probably scrap this altogether and extendAIELiveRegs.cpp

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

llvm/test/CodeGen/AIE/aie2/GlobalISel/post-select-licm-reserved.mir

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

Currently this only support sinking into the exit block. Hoisting to the pre-header will come in a future commit. Ultimately, this pass could become generic code, as this uses a generic target hook to determine what is a "safe" reserved register.

gbossu · 2024-08-05T08:55:25Z

I think I have addressed all the comments now :)

llvm/lib/Target/AIE/ReservedRegsLICM.cpp

martien-de-jong

I'm not sure whether you check for dedicated exit to hoist into. For the rest it looks fine.

The pass will identify phys regs that are defined once in the loop and are not livein.

andcarminati

LGTM. Nice work!

gbossu requested review from abhinay-anubola, abnikant, andcarminati, khallouh, konstantinschwarz, martien-de-jong, SagarMaheshwari99 and stephenneuendorffer as code owners July 22, 2024 09:55

andcarminati reviewed Jul 26, 2024

View reviewed changes

llvm/lib/Target/AIE/ReservedRegsLICM.cpp Outdated Show resolved Hide resolved

martien-de-jong reviewed Jul 30, 2024

View reviewed changes