[llvm] [RFC][SelectionDAG] Not issue TRAP node if naked function (PR #132147)

[via llvm-commits]

https://github.com/yonghong-song updated https://github.com/llvm/llvm-project/pull/132147

>From 63e3e60f1531a62e05ed4ec4c6199b34923bb04c Mon Sep 17 00:00:00 2001
From: Yonghong Song <yonghong.song at linux.dev>
Date: Wed, 19 Mar 2025 14:25:16 -0700
Subject: [PATCH] [RFC][SelectionDAG] Not issue TRAP node if naked function

In [1], Nikita Popov suggested that during lowering 'unreachable' insn
should not generate extra code for naked functions, and this applies
to all architectures. Note that for naked functions, 'unreachable'
insn is necessary in IR since the basic block needs a terminator to end.

This patch checked whether a function is naked function or not.
If it is a naked function, 'unreachable' insn will not generate
ISD::TRAP.

I put the RFC in the subject as I am not sure how to find a test
which can validate the change in FastISel.cpp. Any advice is welcome.

  [1] https://github.com/llvm/llvm-project/pull/131731
---
 llvm/lib/CodeGen/SelectionDAG/FastISel.cpp            |  3 +++
 llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp |  3 +++
 .../WebAssembly/naked-fn-with-frame-pointer.ll        |  2 --
 .../CodeGen/X86/naked-fn-with-unreachable-trap.ll     | 11 +++++++++++
 4 files changed, 17 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/naked-fn-with-unreachable-trap.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
index 33dd039c4ab2a..1bbb704fc0774 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
@@ -1862,6 +1862,9 @@ bool FastISel::selectOperator(const User *I, unsigned Opcode) {
           return true;
       }
 
+      if (cast<Instruction>(I)->getFunction()->hasFnAttribute(Attribute::Naked))
+        return true;
+
       return fastEmit_(MVT::Other, MVT::Other, ISD::TRAP) != 0;
     }
     return true;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 14bb1d943d2d6..6480135a3b929 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3514,6 +3514,9 @@ void SelectionDAGBuilder::visitUnreachable(const UnreachableInst &I) {
       return;
   }
 
+  if (I.getFunction()->hasFnAttribute(Attribute::Naked))
+    return;
+
   DAG.setRoot(DAG.getNode(ISD::TRAP, getCurSDLoc(), MVT::Other, DAG.getRoot()));
 }
 
diff --git a/llvm/test/CodeGen/WebAssembly/naked-fn-with-frame-pointer.ll b/llvm/test/CodeGen/WebAssembly/naked-fn-with-frame-pointer.ll
index fcd42e8cbfb9f..d50cf0b201229 100644
--- a/llvm/test/CodeGen/WebAssembly/naked-fn-with-frame-pointer.ll
+++ b/llvm/test/CodeGen/WebAssembly/naked-fn-with-frame-pointer.ll
@@ -9,13 +9,11 @@ define dso_local void @naked() naked "frame-pointer"="all" {
 ; CHECK-32:         .functype naked () -> ()
 ; CHECK-32-NEXT:  # %bb.0:
 ; CHECK-32-NEXT:    call main
-; CHECK-32-NEXT:    unreachable
 ;
 ; CHECK-64-LABEL: naked:
 ; CHECK-64:         .functype naked () -> ()
 ; CHECK-64-NEXT:  # %bb.0:
 ; CHECK-64-NEXT:    call main
-; CHECK-64-NEXT:    unreachable
   call void @main()
   unreachable
 }
diff --git a/llvm/test/CodeGen/X86/naked-fn-with-unreachable-trap.ll b/llvm/test/CodeGen/X86/naked-fn-with-unreachable-trap.ll
new file mode 100644
index 0000000000000..94274fcb1c160
--- /dev/null
+++ b/llvm/test/CodeGen/X86/naked-fn-with-unreachable-trap.ll
@@ -0,0 +1,11 @@
+; RUN: llc -o - %s -mtriple=x86_64-linux-gnu -trap-unreachable | FileCheck %s
+; RUN: llc -o - %s -mtriple=x86_64-linux-gnu -trap-unreachable -fast-isel | FileCheck %s
+
+define dso_local void @foo() #0 {
+entry:
+  tail call void asm sideeffect "movl 3,%eax", "~{dirflag},~{fpsr},~{flags}"()
+  unreachable
+}
+; CHECK-NOT: ud2
+
+attributes #0 = { naked }