1 files changed, 43 insertions, 2 deletions
diff --git a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
index 8d889a73f9f..4b85664673f 100644
--- a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
+++ b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
@@ -11,15 +11,18 @@ define void @f0() "patchable-function-entry"="0" "branch-target-enforcement" {
 }
 
 ;; -fpatchable-function-entry=1 -mbranch-protection=bti
+;; For M=0, place the label .Lpatch0 after the initial BTI.
 define void @f1() "patchable-function-entry"="1" "branch-target-enforcement" {
 ; CHECK-LABEL: f1:
 ; CHECK-NEXT: .Lfunc_begin1:
-; CHECK:       hint #34
+; CHECK:      // %bb.0:
+; CHECK-NEXT:  hint #34
+; CHECK-NEXT: .Lpatch0:
 ; CHECK-NEXT:  nop
 ; CHECK-NEXT:  ret
 ; CHECK:      .section __patchable_function_entries,"awo",@progbits,f1,unique,0
 ; CHECK-NEXT: .p2align 3
-; CHECK-NEXT: .xword .Lfunc_begin1
+; CHECK-NEXT: .xword .Lpatch0
   ret void
 }
 
@@ -41,3 +44,41 @@ define void @f2_1() "patchable-function-entry"="1" "patchable-function-prefix"="
 ; CHECK-NEXT: .xword .Ltmp0
   ret void
 }
+
+;; -fpatchable-function-entry=1 -mbranch-protection=bti
+;; For M=0, don't create .Lpatch0 if the initial instruction is not BTI,
+;; even if other basic blocks may have BTI.
+define internal void @f1i(i64 %v) "patchable-function-entry"="1" "branch-target-enforcement" {
+; CHECK-LABEL: f1i:
+; CHECK-NEXT: .Lfunc_begin3:
+; CHECK:      // %bb.0:
+; CHECK-NEXT:  nop
+;; Other basic blocks have BTI, but they don't affect our decision to not create .Lpatch0
+; CHECK:      .LBB{{.+}} // %sw.bb1
+; CHECK-NEXT:  hint #36
+; CHECK:      .section __patchable_function_entries,"awo",@progbits,f1,unique,0
+; CHECK-NEXT: .p2align 3
+; CHECK-NEXT: .xword .Lfunc_begin3
+entry:
+  switch i64 %v, label %sw.bb0 [
+    i64 1, label %sw.bb1
+    i64 2, label %sw.bb2
+    i64 3, label %sw.bb3
+    i64 4, label %sw.bb4
+  ]
+sw.bb0:
+  call void asm sideeffect "", ""()
+  ret void
+sw.bb1:
+  call void asm sideeffect "", ""()
+  ret void
+sw.bb2:
+  call void asm sideeffect "", ""()
+  ret void
+sw.bb3:
+  call void asm sideeffect "", ""()
+  ret void
+sw.bb4:
+  call void asm sideeffect "", ""()
+  ret void
+}