[X86][AMX] Try to hoist AMX shapes' def

[pengfei]

Summary

We request no intersections between AMX instructions and their shapes'
def when we insert ldtilecfg. However, this is not always ture resulting
from not only users don't follow AMX API model, but also optimizations.

This patch adds a mechanism that tries to hoist AMX shapes' def as well.
It only hoists shapes inside a BB, we can improve it for cases across
BBs in future. Currently, it only hoists shapes of which all sources' def
above the first AMX instruction. We can improve for the case that only
source that moves an immediate value to a register below AMX instruction.

[pengfei]

Reset ShapeBBs[MBB] once we hoist any shapes in MBB.

[pengfei]

Always insert before AMX instruction.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

199–201

ShapeBBs[MBB] may both contains the Row_def and Col_def, and they may be in different BBs, The Pos of them in their BB my be equal, So use lower_bound here seems not good. Because it may duplicated insert X_def at line 201.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

264

Looks strange here: All BB followed a AMX BB will be mark HasAMXRegLiveIn, Even BBs after the last BB of using AMX?

292

I think "across BBs" is most common case:

AMX BB0 --> shape def --> AMX BB1

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

199–201

ShapeBBs[MBB] contains Shapes only defined in the same MBB. We don't distinguish Row and Col here.

264

It's true from perspective of the live range of TMM registers. Though might be not precise. (The live range actually ends after function call). We simply use such conception here to check if there's any shape def BBs have AMX registers live in, i.e. the shapes are interact with AMX instructions.

292

We can improve it in follow up patches :)

[pengfei]

Minor refactor.

[xiangzhangllvm]

+1 first

[pengfei]

I happened to commit this patch with another one. Since Xiang had +1 for it, @LuoYuanke , do you think if I need to revert it or not?

[hctim]

Looks like this patch broke the MSan buildbots: https://lab.llvm.org/buildbot/#/builders/5/builds/6967/steps/9/logs/stdio

Reproing MSan is a little tricky as it can require a multistage compilation, so you might find the instructions at https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild helpful.

[pengfei]

In D101067#2712921, @hctim wrote:

Looks like this patch broke the MSan buildbots: https://lab.llvm.org/buildbot/#/builders/5/builds/6967/steps/9/logs/stdio

Reproing MSan is a little tricky as it can require a multistage compilation, so you might find the instructions at https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild helpful.

Thanks @hctim for reverting this and the instructions. I can reproduce it locally.

[pengfei]

Reapply the minor refactor.
Fix MSan build failures.

[pengfei]

Ping~

llvm/lib/Target/X86/X86PreTileConfig.cpp

struct MIRef {
   // A virtual position for instruction that will be inserted after MI.
   size_t Pos = 0;
   MIRef() = default;
   MIRef(MachineBasicBlock *MBB) : MBB(MBB) {
     for (auto I = MBB->begin(), E = MBB->end(); I != E && I->isPHI();
          ++I, ++Pos)
       MI = &*I;
   }
+  MIRef(MachineInstr *MI)
+      : MI(MI), MBB(MI->getParent()),
+        Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {}
   MIRef(MachineInstr *MI, MachineBasicBlock *MBB)
       : MI(MI), MBB(MBB),
         Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {}
   MIRef(MachineInstr *MI, MachineBasicBlock *MBB, size_t Pos)
       : MI(MI), MBB(MBB), Pos(Pos) {}
   operator bool() const { return MBB != nullptr; }
   bool operator==(const MIRef &RHS) const {
     return MI == RHS.MI && MBB == RHS.MBB;
   }
+  bool operator!=(const MIRef &RHS) const { return !(*this == RHS); }
   bool operator<(const MIRef &RHS) const {
     return (!MBB && RHS.MBB) || (MBB == RHS.MBB && Pos < RHS.Pos);
   }
   bool operator>(const MIRef &RHS) const {
     return (!RHS.MBB && MBB) || (MBB == RHS.MBB && Pos > RHS.Pos);
   }
 };
 
 struct BBInfo {
   MIRef FirstAMX;
   MIRef LastCall;
-  MIRef LastShape;
+  bool HasAMXRegLiveIn = false;
   bool TileCfgForbidden = false;
   bool NeedTileCfgLiveIn = false;
 };
 
 class X86PreTileConfig : public MachineFunctionPass {
   MachineRegisterInfo *MRI;
   const MachineLoopInfo *MLI;
   SmallSet<MachineInstr *, 8> DefVisited;
-  SmallSet<MachineBasicBlock *, 8> ShapeBBs;
   DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo;
+  DenseMap<MachineBasicBlock *, SmallVector<MIRef, 8>> ShapeBBs;
 
   /// Check if the callee will clobber AMX registers.
   bool isDestructiveCall(MachineInstr &MI, BitVector UsableRegs) {
     auto Iter = llvm::find_if(
         MI.operands(), [](MachineOperand &MO) { return MO.isRegMask(); });
     if (Iter == MI.operands_end())
       return false;
     UsableRegs.clearBitsInMask(Iter->getRegMask());
class X86PreTileConfig : public MachineFunctionPass {
   bool isLoopBackEdge(MachineBasicBlock *Header, MachineBasicBlock *Bottom) {
     return MLI->isLoopHeader(Header) &&
            MLI->getLoopFor(Header)->getBottomBlock() == Bottom;
   }
 
   /// Collect the shape def information for later use.
   void collectShapeInfo(MachineInstr &MI);
 
+  /// Try to hoist shapes definded below AMX instructions.
+  bool hoistShapesInBB(MachineBasicBlock *MBB) {
+    auto FirstShapeBelowAMX =
+        llvm::lower_bound(ShapeBBs[MBB], BBVisitedInfo[MBB].FirstAMX);
+    auto InsertPoint = BBVisitedInfo[MBB].FirstAMX.MI->getIterator();
+    for (auto I = FirstShapeBelowAMX, E = ShapeBBs[MBB].end(); I != E; ++I) {
+      // Do not hoist instructions that access memory.
+      if (I->MI->mayLoadOrStore())
+        return false;
+      for (auto &MO : I->MI->operands()) {
+        if (MO.isDef())
+          continue;
+        // Do not hoist instructions if the sources' def under AMX instruction.
+        // TODO: We can handle isMoveImmediate MI here.
+        if (MO.isReg() &&
+            MIRef(MRI->getVRegDef(MO.getReg())) > BBVisitedInfo[MBB].FirstAMX)
+          return false;
+        // TODO: Maybe need more checks here.
+      }
+      InsertPoint = MBB->insert(InsertPoint, I->MI->removeFromParent());
+    }
+    // We only need to mark the last shape in the BB now.
+    ShapeBBs[MBB].clear();
+    ShapeBBs[MBB].push_back(MIRef(&*InsertPoint, MBB));
+    return true;
+  }
+
 public:
   X86PreTileConfig() : MachineFunctionPass(ID) {}
 
   /// Return the pass name.
   StringRef getPassName() const override {
     return "Tile Register Pre-configure";
   }
 
INITIALIZE_PASS_BEGIN(X86PreTileConfig, "tilepreconfig",
                       "Tile Register Pre-configure", false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_END(X86PreTileConfig, "tilepreconfig",
                     "Tile Register Pre-configure", false, false)
 
 void X86PreTileConfig::collectShapeInfo(MachineInstr &MI) {
   auto RecordShape = [&](MachineInstr *MI, MachineBasicBlock *MBB) {
     MIRef MIR(MI, MBB);
-    if (BBVisitedInfo[MBB].LastShape < MIR)
-      BBVisitedInfo[MBB].LastShape = MIR;
-    ShapeBBs.insert(MBB);
+    auto I = llvm::lower_bound(ShapeBBs[MBB], MIR);
+    if (*I != MIR)
+      ShapeBBs[MBB].insert(I, MIR);
   };
 
   SmallVector<Register, 8> WorkList(
       {MI.getOperand(1).getReg(), MI.getOperand(2).getReg()});
   while (!WorkList.empty()) {
     Register R = WorkList.pop_back_val();
     MachineInstr *DefMI = MRI->getVRegDef(R);
     MachineBasicBlock *DefMBB = DefMI->getParent();
for (auto &MI : MBB) {
       }
     }
     if (BBVisitedInfo[&MBB].NeedTileCfgLiveIn) {
       if (&MBB == &MF.front())
         CfgNeedInsert.insert(MIRef(&MBB));
       else
         CfgLiveInBBs.push_back(&MBB);
     }
+    if (BBVisitedInfo[&MBB].FirstAMX || BBVisitedInfo[&MBB].HasAMXRegLiveIn)
+      for (auto *Succ : MBB.successors())
+        if (!isLoopBackEdge(Succ, &MBB))
+          BBVisitedInfo[Succ].HasAMXRegLiveIn = true;
   }
 
   // Update NeedTileCfgLiveIn for predecessors.
   while (!CfgLiveInBBs.empty()) {
     MachineBasicBlock *MBB = CfgLiveInBBs.pop_back_val();
     for (auto *Pred : MBB->predecessors()) {
       if (BBVisitedInfo[Pred].LastCall) {
         CfgNeedInsert.insert(BBVisitedInfo[Pred].LastCall);
       } else if (!BBVisitedInfo[Pred].NeedTileCfgLiveIn) {
         BBVisitedInfo[Pred].NeedTileCfgLiveIn = true;
         if (Pred == &MF.front())
           CfgNeedInsert.insert(MIRef(Pred));
         else
           CfgLiveInBBs.push_back(Pred);
       }
     }
   }
 
   // There's no AMX instruction if we didn't find a tile config live in point.
   if (CfgNeedInsert.empty())
     return false;
 
   // Avoid to insert ldtilecfg before any shape defs.
-  SmallVector<MachineBasicBlock *, 8> WorkList(
-      make_range(ShapeBBs.begin(), ShapeBBs.end()));
+  SmallVector<MachineBasicBlock *, 8> WorkList;
+  for (auto &I : ShapeBBs) {
+    // TODO: We can hoist shapes across BBs here.
+    if (BBVisitedInfo[I.first].HasAMXRegLiveIn)
+      REPORT_CONFIG_FAIL
+    if (BBVisitedInfo[I.first].FirstAMX &&
+        BBVisitedInfo[I.first].FirstAMX < ShapeBBs[I.first].back() &&
+        !hoistShapesInBB(I.first))
+      REPORT_CONFIG_FAIL
+    WorkList.push_back(I.first);
+  }
   while (!WorkList.empty()) {
     MachineBasicBlock *MBB = WorkList.pop_back_val();
     for (auto *Pred : MBB->predecessors()) {
       if (!BBVisitedInfo[Pred].TileCfgForbidden && !isLoopBackEdge(MBB, Pred)) {
         BBVisitedInfo[Pred].TileCfgForbidden = true;
         WorkList.push_back(Pred);
       }
     }
while (!WorkList.empty()) {
       MIRef I = WorkList.pop_back_val();
       if (!VisitedOrInserted.count(I)) {
         if (!BBVisitedInfo[I.MBB].TileCfgForbidden) {
           // If the BB is all shapes reachable, stop sink and try to insert.
           InsertPoints.insert(I);
         } else {
           // Avoid the BB to be multi visited.
           VisitedOrInserted.insert(I);
-          // We cannot sink it across any AMX instruction.
-          if (BBVisitedInfo[I.MBB].FirstAMX)
-            REPORT_CONFIG_FAIL;
           // Sink the inserting point along the chain with NeedTileCfgLiveIn =
           // true when MBB isn't all shapes reachable.
           for (auto *Succ : I.MBB->successors())
             if (BBVisitedInfo[Succ].NeedTileCfgLiveIn)
               WorkList.push_back(MIRef(Succ));
         }
       }
     }
 
     // A given point might be forked due to shape conditions are not met.
     for (MIRef I : InsertPoints) {
-      // Even MBB is all shapes reachable, we still need to check if there's
-      // AMX that intersects with shapes in the same MBB.
-      if (BBVisitedInfo[I.MBB].FirstAMX &&
-          BBVisitedInfo[I.MBB].FirstAMX < BBVisitedInfo[I.MBB].LastShape)
-        REPORT_CONFIG_FAIL;
       // Make sure we insert ldtilecfg after the last shape def in MBB.
-      if (I < BBVisitedInfo[I.MBB].LastShape)
-        I = BBVisitedInfo[I.MBB].LastShape;
+      if (ShapeBBs.count(I.MBB) && I < ShapeBBs[I.MBB].back())
+        I = ShapeBBs[I.MBB].back();
       // There're chances the MBB is sunk more than once. Record it to avoid
       // multi insert.
       if (VisitedOrInserted.insert(I).second) {
         auto II = I.MI ? I.MI->getIterator() : I.MBB->instr_begin();
         addFrameReference(BuildMI(*I.MBB, ++II, DL, TII->get(X86::LDTILECFG)),
                           SS);
       }
     }

llvm/test/CodeGen/X86/AMX/amx-sched.ll

 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx512f -mcpu=skx -verify-machineinstrs | FileCheck %s
 
 define <256 x i32> @test_shape_sched(i16 %m, i16 %n, i16 %k, <256 x i32> %c, <256 x i32> %a, <256 x i32> %b) nounwind {
 ; Just to make sure shape def is not scheduled across ldtilecfg.
+; CHECK-LABEL: test_shape_sched:
 ; CHECK:    ldtilecfg
 ; CHECK-NOT: movw
   %c1 = bitcast <256 x i32> %c to x86_amx
   %a1 = bitcast <256 x i32> %a to x86_amx
   %b1 = bitcast <256 x i32> %b to x86_amx
   %t = call x86_amx @llvm.x86.tdpbssd.internal(i16 %m, i16 %n, i16 %k, x86_amx %c1, x86_amx %a1, x86_amx %b1)
   %res = bitcast x86_amx %t to <256 x i32>
   ret <256 x i32> %res
 }
 
+define <256 x i32> @test_shape_sched2(i16 %m, i16 %n, i16 %k, i8* %c, i8* %a, i8* %b) nounwind {
+; Just to make sure shape def is not scheduled across ldtilecfg.
+; CHECK-LABEL: test_shape_sched2:
+; CHECK:    ldtilecfg
+; CHECK-NOT: movw
+  %aa = lshr i16 %k, 2
+  %c1 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %m, i16 %n, i8* %c, i64 64)
+  %a1 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %m, i16 %k, i8* %a, i64 64)
+  %b1 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %aa, i16 %n, i8* %b, i64 64)
+  %t = call x86_amx @llvm.x86.tdpbssd.internal(i16 %m, i16 %n, i16 %k, x86_amx %c1, x86_amx %a1, x86_amx %b1)
+  %res = bitcast x86_amx %t to <256 x i32>
+  ret <256 x i32> %res
+}
 
+declare x86_amx @llvm.x86.tileloadd64.internal(i16, i16, i8*, i64)
 declare x86_amx @llvm.x86.tdpbssd.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)