[llvm-branch-commits] [X86] Finish AMX-RegAlloc Related Pipeline Overrides (PR #234079)

[via llvm-branch-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

Author: Aiden Grossman (boomanaiden154)

<details>
<summary>Changes</summary>

The modification points are now available so finish up this part of the
pipeline, just matching the behavior from the X86 TargetPassConfig for
the LegacyPM.


---
Full diff: https://github.com/llvm/llvm-project/pull/234079.diff


2 Files Affected:

- (modified) llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp (+29-4) 
- (modified) llvm/test/CodeGen/X86/llc-pipeline-npm.ll (+2) 


``````````diff
diff --git a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
index 058fc04bb9f59..0bbebcdbf9b9f 100644
--- a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
+++ b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
@@ -28,11 +28,13 @@
 #include "llvm/CodeGen/JMCInstrumenter.h"
 #include "llvm/CodeGen/KCFI.h"
 #include "llvm/CodeGen/MachineCombiner.h"
+#include "llvm/CodeGen/RegAllocGreedyPass.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/Passes/CodeGenPassBuilder.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Transforms/CFGuard.h"
 
 using namespace llvm;
@@ -66,14 +68,13 @@ class X86CodeGenPassBuilder : public CodeGenPassBuilder {
   void addILPOpts(PassManagerWrapper &PMW) override;
   void addMachineSSAOptimization(PassManagerWrapper &PMW) override;
   void addPreRegAlloc(PassManagerWrapper &PMW) override;
-  // TODO(boomanaiden154): We need to add addPostFastRegAllocRewrite here once
-  // it is available to support AMX.
+  void addPostFastRegAllocRewrite(PassManagerWrapper &PMW) override;
   void addPostRegAlloc(PassManagerWrapper &PMW) override;
   void addPreSched2(PassManagerWrapper &PMW) override;
   void addPreEmitPass(PassManagerWrapper &PMW) override;
   void addPreEmitPass2(PassManagerWrapper &PMW) override;
-  // TODO(boomanaiden154): We need to add addRegAssignAndRewriteOptimized here
-  // once it is available to support AMX.
+  Expected<bool>
+  addRegAssignAndRewriteOptimized(PassManagerWrapper &PMW) override;
   void addAsmPrinterBegin(PassManagerWrapper &PMW) override;
   void addAsmPrinter(PassManagerWrapper &PMW) override;
   void addAsmPrinterEnd(PassManagerWrapper &PMW) override;
@@ -199,6 +200,11 @@ void X86CodeGenPassBuilder::addPreRegAlloc(PassManagerWrapper &PMW) {
     addMachineFunctionPass(X86FastPreTileConfigPass(), PMW);
 }
 
+void X86CodeGenPassBuilder::addPostFastRegAllocRewrite(
+    PassManagerWrapper &PMW) {
+  addMachineFunctionPass(X86FastTileConfigPass(), PMW);
+}
+
 void X86CodeGenPassBuilder::addPostRegAlloc(PassManagerWrapper &PMW) {
   addMachineFunctionPass(X86LowerTileCopyPass(), PMW);
   addMachineFunctionPass(X86FPStackifierPass(), PMW);
@@ -309,6 +315,25 @@ void X86CodeGenPassBuilder::addPreEmitPass2(PassManagerWrapper &PMW) {
   }
 }
 
+static bool onlyAllocateTileRegisters(const TargetRegisterInfo &TRI,
+                                      const MachineRegisterInfo &MRI,
+                                      const Register Reg) {
+  const TargetRegisterClass *RC = MRI.getRegClass(Reg);
+  return static_cast<const X86RegisterInfo &>(TRI).isTileRegisterClass(RC);
+}
+
+Expected<bool> X86CodeGenPassBuilder::addRegAssignAndRewriteOptimized(
+    PassManagerWrapper &PMW) {
+  // Don't support tile RA when RA is specified by command line "-regalloc".
+  if (Opt.RegAlloc == RegAllocType::Default && getTM().getCLOpts().tile_ra) {
+    // Allocate tile register first.
+    addMachineFunctionPass(
+        RAGreedyPass({RAGreedyPass::Options(onlyAllocateTileRegisters)}), PMW);
+    addMachineFunctionPass(X86TileConfigPass(), PMW);
+  }
+  return Base::addRegAssignAndRewriteOptimized(PMW);
+}
+
 void X86CodeGenPassBuilder::addAsmPrinterBegin(PassManagerWrapper &PMW) {
   addModulePass(X86AsmPrinterBeginPass(), PMW, /*Force=*/true);
 }
diff --git a/llvm/test/CodeGen/X86/llc-pipeline-npm.ll b/llvm/test/CodeGen/X86/llc-pipeline-npm.ll
index 7138ab5e8c671..0c82a039683e5 100644
--- a/llvm/test/CodeGen/X86/llc-pipeline-npm.ll
+++ b/llvm/test/CodeGen/X86/llc-pipeline-npm.ll
@@ -50,6 +50,7 @@
 ; O0-NEXT:     x86-fast-pre-tile-config
 ; O0-NEXT:     phi-node-elimination
 ; O0-NEXT:     regallocfast
+; O0-NEXT:     x86-fast-tile-config
 ; O0-NEXT:     x86-lower-tile-copy
 ; O0-NEXT:     x86-fp-stackifier
 ; O0-NEXT:     remove-redundant-debug-values
@@ -253,6 +254,7 @@
 ; O0-WINDOWS-NEXT:     x86-fast-pre-tile-config
 ; O0-WINDOWS-NEXT:     phi-node-elimination
 ; O0-WINDOWS-NEXT:     regallocfast
+; O0-WINDOWS-NEXT:     x86-fast-tile-config
 ; O0-WINDOWS-NEXT:     x86-lower-tile-copy
 ; O0-WINDOWS-NEXT:     x86-fp-stackifier
 ; O0-WINDOWS-NEXT:     remove-redundant-debug-values

``````````

</details>


https://github.com/llvm/llvm-project/pull/234079