[llvm] [BPF] Remove 'may_goto 0' instructions (PR #123482)

[via llvm-commits]

https://github.com/yonghong-song created https://github.com/llvm/llvm-project/pull/123482

Emil Tsalapatis from Meta reported such a case where 'may_goto 0' insn is generated by clang compiler. But 'may_goto 0' insn is actually a no-op so it makes sense to remove that in llvm. The patch is also able to handle the following code pattern
```
   ...
   may_goto 2
   may_goto 1
   may_goto 0
   ...
```
where three may_goto insns can all be removed.

>From cb821087be9072b63b2bfa3ef9c5e20696dcbf4c Mon Sep 17 00:00:00 2001
From: Yonghong Song <yonghong.song at linux.dev>
Date: Mon, 13 Jan 2025 13:31:24 -0800
Subject: [PATCH] [BPF] Remove 'may_goto 0' instructions

Emil Tsalapatis from Meta reported such a case where 'may_goto 0' insn is
generated by clang compiler. But 'may_goto 0' insn is actually a no-op so
it makes sense to remove that in llvm. The patch is also able to handle
the following code pattern
   ...
   may_goto 2
   may_goto 1
   may_goto 0
   ...
where three may_goto insns can all be removed.
---
 llvm/lib/Target/BPF/BPFMIPeephole.cpp | 62 +++++++++++++++++++++++++++
 llvm/test/CodeGen/BPF/may_goto.ll     | 27 ++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 llvm/test/CodeGen/BPF/may_goto.ll

diff --git a/llvm/lib/Target/BPF/BPFMIPeephole.cpp b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
index 2b17f2aaefe2bf..eda7a0f0acd363 100644
--- a/llvm/lib/Target/BPF/BPFMIPeephole.cpp
+++ b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
@@ -24,6 +24,7 @@
 #include "BPFInstrInfo.h"
 #include "BPFTargetMachine.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
@@ -322,6 +323,7 @@ struct BPFMIPreEmitPeephole : public MachineFunctionPass {
   bool eliminateRedundantMov();
   bool adjustBranch();
   bool insertMissingCallerSavedSpills();
+  bool removeMayGotoZero();
 
 public:
 
@@ -337,6 +339,7 @@ struct BPFMIPreEmitPeephole : public MachineFunctionPass {
     if (SupportGotol)
       Changed = adjustBranch() || Changed;
     Changed |= insertMissingCallerSavedSpills();
+    Changed |= removeMayGotoZero();
     return Changed;
   }
 };
@@ -682,6 +685,65 @@ bool BPFMIPreEmitPeephole::insertMissingCallerSavedSpills() {
   return Changed;
 }
 
+bool BPFMIPreEmitPeephole::removeMayGotoZero() {
+  bool Changed = false;
+  MachineBasicBlock *Prev_MBB, *Curr_MBB = nullptr;
+  MachineBasicBlock *ToErase = nullptr;
+
+  for (MachineBasicBlock &MBB : reverse(*MF)) {
+    if (ToErase) {
+      ToErase->eraseFromParent();
+      ToErase = nullptr;
+    }
+
+    Prev_MBB = Curr_MBB;
+    Curr_MBB = &MBB;
+    if (Prev_MBB == nullptr)
+      continue;
+
+    MachineInstr &MI = MBB.back();
+    if (!MI.isInlineAsm())
+      continue;
+
+    const char *AsmStr = MI.getOperand(0).getSymbolName();
+    SmallVector<StringRef, 4> AsmPieces;
+    SplitString(AsmStr, AsmPieces, ";\n");
+
+    // Do not support multiple insns in one inline asm.
+    if (AsmPieces.size() != 1)
+      continue;
+
+    // The asm insn must be a may_goto insn.
+    SmallVector<StringRef, 4> AsmOpPieces;
+    SplitString(AsmPieces[0], AsmOpPieces, " ");
+    if (AsmOpPieces[0] != "may_goto")
+      continue;
+
+    // Get the may_goto branch target.
+    MachineOperand &MO = MI.getOperand(InlineAsm::MIOp_FirstOperand + 1);
+    if (MO.getMBB() != Prev_MBB)
+      continue;
+
+    Changed = true;
+    if (MBB.begin() == MI) {
+      // Single 'may_goto' insn in the same basic block.
+      ToErase = Curr_MBB;
+      Curr_MBB->removeSuccessor(Prev_MBB);
+      for (MachineBasicBlock *Pred : Curr_MBB->predecessors())
+        Pred->replaceSuccessor(Curr_MBB, Prev_MBB);
+      Curr_MBB = Prev_MBB;
+    } else {
+      // Remove 'may_goto' insn.
+      MI.eraseFromParent();
+    }
+  }
+
+  if (ToErase)
+    ToErase->eraseFromParent();
+
+  return Changed;
+}
+
 } // end default namespace
 
 INITIALIZE_PASS(BPFMIPreEmitPeephole, "bpf-mi-pemit-peephole",
diff --git a/llvm/test/CodeGen/BPF/may_goto.ll b/llvm/test/CodeGen/BPF/may_goto.ll
new file mode 100644
index 00000000000000..a47716509c7aba
--- /dev/null
+++ b/llvm/test/CodeGen/BPF/may_goto.ll
@@ -0,0 +1,27 @@
+; RUN: llc -mtriple=bpfel -mcpu=v3 -filetype=obj -o - %s | llvm-objdump --no-show-raw-insn -d - | FileCheck %s
+
+ at j = dso_local local_unnamed_addr global i32 0, align 4
+
+define dso_local noundef i32 @foo() local_unnamed_addr {
+entry:
+  callbr void asm sideeffect "may_goto ${0:l}", "!i"()
+          to label %for.body [label %for.cond.cleanup]
+
+for.cond.cleanup:                                 ; preds = %for.body.2, %for.body.2, %for.body.1, %for.body, %entry
+  ret i32 0
+
+for.body:                                         ; preds = %entry
+  callbr void asm sideeffect "may_goto ${0:l}", "!i"()
+          to label %for.body.1 [label %for.cond.cleanup]
+
+for.body.1:                                       ; preds = %for.body
+  callbr void asm sideeffect "may_goto ${0:l}", "!i"()
+          to label %for.body.2 [label %for.cond.cleanup]
+
+for.body.2:                                       ; preds = %for.body.1
+  callbr void asm sideeffect "may_goto ${0:l}", "!i"()
+          to label %for.cond.cleanup [label %for.cond.cleanup]
+}
+
+; CHECK:       0:       w0 = 0x0
+; CHECK-NEXT:  1:       exit