[llvm] [AMDGPU][GISel] Notify GISel observers when splitting trap blocks (PR #219128)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Keshav Vinayak Jha (keshavvinayak01)

<details>
<summary>Changes</summary>

`legalizeTrapEndpgm` can split a MBB and move the instructions following a trap. Since an instruction's parent block is part of its GISel CSE profile, notify active observers around the moved instructions so they are re-profiled with the new parent block.

CSE consistency gap exposed by #<!-- -->217258.

Assisted-by: Codex

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


3 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp (+21-7) 
- (modified) llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h (+2-4) 
- (modified) llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-trap.mir (+4) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 7c1a26f761c96..a794c41c1c86b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -24,6 +24,8 @@
 #include "SIRegisterInfo.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
@@ -2412,7 +2414,7 @@ bool AMDGPULegalizerInfo::legalizeCustom(
   case TargetOpcode::G_SET_FPENV:
     return legalizeSetFPEnv(MI, MRI, B);
   case TargetOpcode::G_TRAP:
-    return legalizeTrap(MI, MRI, B);
+    return legalizeTrap(Helper, MI);
   case TargetOpcode::G_DEBUGTRAP:
     return legalizeDebugTrap(MI, MRI, B);
   default:
@@ -7746,19 +7748,22 @@ bool AMDGPULegalizerInfo::legalizeSBufferPrefetch(LegalizerHelper &Helper,
 }
 
 // TODO: Move to selection
-bool AMDGPULegalizerInfo::legalizeTrap(MachineInstr &MI,
-                                       MachineRegisterInfo &MRI,
-                                       MachineIRBuilder &B) const {
+bool AMDGPULegalizerInfo::legalizeTrap(LegalizerHelper &Helper,
+                                       MachineInstr &MI) const {
+  MachineIRBuilder &B = Helper.MIRBuilder;
+  MachineRegisterInfo &MRI = *B.getMRI();
   if (!ST.hasTrapHandler() ||
       ST.getTrapHandlerAbi() != GCNSubtarget::TrapHandlerAbi::AMDHSA)
-    return legalizeTrapEndpgm(MI, MRI, B);
+    return legalizeTrapEndpgm(Helper, MI);
 
   return ST.supportsGetDoorbellID() ?
          legalizeTrapHsa(MI, MRI, B) : legalizeTrapHsaQueuePtr(MI, MRI, B);
 }
 
-bool AMDGPULegalizerInfo::legalizeTrapEndpgm(
-    MachineInstr &MI, MachineRegisterInfo &MRI, MachineIRBuilder &B) const {
+bool AMDGPULegalizerInfo::legalizeTrapEndpgm(LegalizerHelper &Helper,
+                                             MachineInstr &MI) const {
+  MachineIRBuilder &B = Helper.MIRBuilder;
+  GISelChangeObserver &Observer = Helper.Observer;
   const DebugLoc &DL = MI.getDebugLoc();
   MachineBasicBlock &BB = B.getMBB();
   MachineFunction *MF = BB.getParent();
@@ -7773,7 +7778,16 @@ bool AMDGPULegalizerInfo::legalizeTrapEndpgm(
   // We need a block split to make the real endpgm a terminator. We also don't
   // want to break phis in successor blocks, so we can't just delete to the
   // end of the block.
+  // An instruction's parent block is part of its CSE profile, so notify
+  // observers about the instructions moved by the split.
+  SmallVector<MachineInstr *, 8> MovedInstrs;
+  for (auto I = std::next(MI.getIterator()), E = BB.end(); I != E; ++I) {
+    Observer.changingInstr(*I);
+    MovedInstrs.push_back(&*I);
+  }
   BB.splitAt(MI, false /*UpdateLiveIns*/);
+  for (MachineInstr *MovedMI : MovedInstrs)
+    Observer.changedInstr(*MovedMI);
   MachineBasicBlock *TrapBB = MF->CreateMachineBasicBlock();
   MF->push_back(TrapBB);
   BuildMI(*TrapBB, TrapBB->end(), DL, B.getTII().get(AMDGPU::S_ENDPGM))
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
index 89819fe990f5a..c68568f775dc7 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
@@ -251,10 +251,8 @@ class AMDGPULegalizerInfo final : public LegalizerInfo {
 
   bool legalizeSBufferPrefetch(LegalizerHelper &Helper, MachineInstr &MI) const;
 
-  bool legalizeTrap(MachineInstr &MI, MachineRegisterInfo &MRI,
-                    MachineIRBuilder &B) const;
-  bool legalizeTrapEndpgm(MachineInstr &MI, MachineRegisterInfo &MRI,
-                          MachineIRBuilder &B) const;
+  bool legalizeTrap(LegalizerHelper &Helper, MachineInstr &MI) const;
+  bool legalizeTrapEndpgm(LegalizerHelper &Helper, MachineInstr &MI) const;
   bool legalizeTrapHsaQueuePtr(MachineInstr &MI, MachineRegisterInfo &MRI,
                                MachineIRBuilder &B) const;
   bool legalizeTrapHsa(MachineInstr &MI, MachineRegisterInfo &MRI,
diff --git a/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-trap.mir b/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-trap.mir
index 80e88e9ed56f0..72567a3783172 100644
--- a/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-trap.mir
+++ b/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-trap.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
 # RUN: llc -mtriple=amdgpu6.00-mesa-mesa3d -run-pass=legalizer -o - %s | FileCheck -check-prefix=GCN %s
+# RUN: %if asserts %{ llc -mtriple=amdgpu6.00-mesa-mesa3d -run-pass=legalizer -enable-cse-in-legalizer=1 -debug-only=cseinfo -filetype=null %s 2>&1 | FileCheck -check-prefix=CSE %s %}
 
 # Check edge cases for trap legalization
 
@@ -34,6 +35,9 @@ body: |
 ---
 name: test_def_fallthrough_after_trap
 body: |
+  ; CSE: CSEInfo::Add MI: %1:_(p1) = G_CONSTANT i64 0
+  ; CSE: CSEInfo::Recording new MI %1:_(p1) = G_CONSTANT i64 0
+  ; CSE: CSEInfo::Recording new MI %1:_(p1) = G_CONSTANT i64 0
   ; GCN-LABEL: name: test_def_fallthrough_after_trap
   ; GCN: bb.0:
   ; GCN-NEXT:   successors: %bb.2(0x40000000), %bb.3(0x40000000)

``````````

</details>


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