[llvm] [AMDGPU] Handled G_UBSANTRAP GlobalIsel (PR #134492)

[via llvm-commits]

https://github.com/amansharma612 created https://github.com/llvm/llvm-project/pull/134492

Fixes #131740

>From 515510f1ef088d2ed972d4b00e5660e207bb55b0 Mon Sep 17 00:00:00 2001
From: Aman Sharma <210100011 at iitb.ac.in>
Date: Sat, 5 Apr 2025 16:02:22 +0530
Subject: [PATCH 1/2] [AMDGPU] Added G_UBSANTRAP GlobalISel

---
 .../lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 27 ++++++++++++++++++-
 llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h  |  5 ++++
 llvm/test/CodeGen/AMDGPU/ubsan_trap.ll        |  7 +++++
 3 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/ubsan_trap.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 275d0193452a5..a317bf884c984 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -2103,7 +2103,7 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
   getActionDefinitionsBuilder({G_MEMCPY, G_MEMCPY_INLINE, G_MEMMOVE, G_MEMSET})
       .lower();
 
-  getActionDefinitionsBuilder({G_TRAP, G_DEBUGTRAP}).custom();
+  getActionDefinitionsBuilder({G_TRAP, G_DEBUGTRAP, G_UBSANTRAP}).custom();
 
   getActionDefinitionsBuilder({G_VASTART, G_VAARG, G_BRJT, G_JUMP_TABLE,
         G_INDEXED_LOAD, G_INDEXED_SEXTLOAD,
@@ -2222,6 +2222,8 @@ bool AMDGPULegalizerInfo::legalizeCustom(
     return legalizeTrap(MI, MRI, B);
   case TargetOpcode::G_DEBUGTRAP:
     return legalizeDebugTrap(MI, MRI, B);
+  case TargetOpcode::G_UBSANTRAP:
+    return legalizeUbsanTrap(MI, MRI, B);
   default:
     return false;
   }
@@ -7045,6 +7047,29 @@ bool AMDGPULegalizerInfo::legalizeDebugTrap(MachineInstr &MI,
   return true;
 }
 
+
+bool AMDGPULegalizerInfo::legalizeUbsanTrap(MachineInstr &MI,
+                                            MachineRegisterInfo &MRI,
+                                            MachineIRBuilder &B) const {
+  // Is non-HSA path or trap-handler disabled? Then, report a warning
+  // accordingly
+  if (!ST.isTrapHandlerEnabled() ||
+      ST.getTrapHandlerAbi() != GCNSubtarget::TrapHandlerAbi::AMDHSA) {
+    DiagnosticInfoUnsupported NoTrap(B.getMF().getFunction(),
+                                     "ubsantrap handler not supported",
+                                     MI.getDebugLoc(), DS_Warning);
+    LLVMContext &Ctx = B.getMF().getFunction().getContext();
+    Ctx.diagnose(NoTrap);
+  } else {
+    // Insert trap instruction
+    B.buildInstr(AMDGPU::S_TRAP)
+        .addImm(static_cast<unsigned>(GCNSubtarget::TrapID::LLVMAMDHSATrap));
+  }
+
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AMDGPULegalizerInfo::legalizeBVHIntersectRayIntrinsic(
     MachineInstr &MI, MachineIRBuilder &B) const {
   MachineRegisterInfo &MRI = *B.getMRI();
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
index 1f4e02b0d600a..5b775b3ef4a10 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
@@ -244,6 +244,11 @@ class AMDGPULegalizerInfo final : public LegalizerInfo {
   bool legalizeDebugTrap(MachineInstr &MI, MachineRegisterInfo &MRI,
                          MachineIRBuilder &B) const;
 
+
+  bool legalizeUbsanTrap(MachineInstr &MI, MachineRegisterInfo &MRI,
+   MachineIRBuilder &B) const;
+
+
   bool legalizeIntrinsic(LegalizerHelper &Helper,
                          MachineInstr &MI) const override;
 };
diff --git a/llvm/test/CodeGen/AMDGPU/ubsan_trap.ll b/llvm/test/CodeGen/AMDGPU/ubsan_trap.ll
new file mode 100644
index 0000000000000..9f961031a1ddb
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/ubsan_trap.ll
@@ -0,0 +1,7 @@
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -global-isel < %s
+; LLVM ERROR: cannot select: G_UBSANTRAP 0 (in function: ubsan_trap)
+
+define void @ubsan_trap() {
+  call void @llvm.ubsantrap(i8 0)
+  ret void
+}
\ No newline at end of file

>From 5d1aa6a393da28c9f124e9dd797296cdc225fec9 Mon Sep 17 00:00:00 2001
From: Aman Sharma <210100011 at iitb.ac.in>
Date: Sat, 5 Apr 2025 16:04:26 +0530
Subject: [PATCH 2/2] added clang-format output

---
 llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 1 -
 llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h   | 4 +---
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index a317bf884c984..dd8a010adb6df 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -7047,7 +7047,6 @@ bool AMDGPULegalizerInfo::legalizeDebugTrap(MachineInstr &MI,
   return true;
 }
 
-
 bool AMDGPULegalizerInfo::legalizeUbsanTrap(MachineInstr &MI,
                                             MachineRegisterInfo &MRI,
                                             MachineIRBuilder &B) const {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
index 5b775b3ef4a10..076a66bb6012f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
@@ -244,10 +244,8 @@ class AMDGPULegalizerInfo final : public LegalizerInfo {
   bool legalizeDebugTrap(MachineInstr &MI, MachineRegisterInfo &MRI,
                          MachineIRBuilder &B) const;
 
-
   bool legalizeUbsanTrap(MachineInstr &MI, MachineRegisterInfo &MRI,
-   MachineIRBuilder &B) const;
-
+                         MachineIRBuilder &B) const;
 
   bool legalizeIntrinsic(LegalizerHelper &Helper,
                          MachineInstr &MI) const override;