[AArch64][SME] Handle SME state around TLS-descriptor calls (#155608)

This patch ensures we switch out of streaming mode before TLS-descriptor
calls. ZA state will also be preserved when using the new SME ABI
lowering (`-aarch64-new-sme-abi`).

Fixes #152165

1 files changed, 159 insertions, 0 deletions

diff --git a/llvm/test/CodeGen/AArch64/sme-dynamic-tls.ll b/llvm/test/CodeGen/AArch64/sme-dynamic-tls.ll
new file mode 100644
index 000000000000..f72ccadea5db
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sme-dynamic-tls.ll
@@ -0,0 +1,159 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -mattr=+sme -aarch64-new-sme-abi -relocation-model=pic < %s | FileCheck %s
+
+@x = external thread_local local_unnamed_addr global i32, align 4
+
+define i32 @load_tls_streaming_compat() nounwind "aarch64_pstate_sm_compatible" {
+; CHECK-LABEL: load_tls_streaming_compat:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp d15, d14, [sp, #-80]! // 16-byte Folded Spill
+; CHECK-NEXT:    stp d13, d12, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d11, d10, [sp, #32] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d9, d8, [sp, #48] // 16-byte Folded Spill
+; CHECK-NEXT:    str x30, [sp, #64] // 8-byte Folded Spill
+; CHECK-NEXT:    mrs x8, SVCR
+; CHECK-NEXT:    tbz w8, #0, .LBB0_2
+; CHECK-NEXT:  // %bb.1: // %entry
+; CHECK-NEXT:    smstop sm
+; CHECK-NEXT:  .LBB0_2: // %entry
+; CHECK-NEXT:    adrp x0, :tlsdesc:x
+; CHECK-NEXT:    ldr x1, [x0, :tlsdesc_lo12:x]
+; CHECK-NEXT:    add x0, x0, :tlsdesc_lo12:x
+; CHECK-NEXT:    .tlsdesccall x
+; CHECK-NEXT:    blr x1
+; CHECK-NEXT:    tbz w8, #0, .LBB0_4
+; CHECK-NEXT:  // %bb.3: // %entry
+; CHECK-NEXT:    smstart sm
+; CHECK-NEXT:  .LBB0_4: // %entry
+; CHECK-NEXT:    mrs x8, TPIDR_EL0
+; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d11, d10, [sp, #32] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr w0, [x8, x0]
+; CHECK-NEXT:    ldp d13, d12, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr x30, [sp, #64] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp d15, d14, [sp], #80 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @x)
+  %1 = load i32, ptr %0, align 4
+  ret i32 %1
+}
+
+define i32 @load_tls_streaming() nounwind "aarch64_pstate_sm_enabled" {
+; CHECK-LABEL: load_tls_streaming:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp d15, d14, [sp, #-80]! // 16-byte Folded Spill
+; CHECK-NEXT:    stp d13, d12, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d11, d10, [sp, #32] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d9, d8, [sp, #48] // 16-byte Folded Spill
+; CHECK-NEXT:    str x30, [sp, #64] // 8-byte Folded Spill
+; CHECK-NEXT:    smstop sm
+; CHECK-NEXT:    adrp x0, :tlsdesc:x
+; CHECK-NEXT:    ldr x1, [x0, :tlsdesc_lo12:x]
+; CHECK-NEXT:    add x0, x0, :tlsdesc_lo12:x
+; CHECK-NEXT:    .tlsdesccall x
+; CHECK-NEXT:    blr x1
+; CHECK-NEXT:    smstart sm
+; CHECK-NEXT:    mrs x8, TPIDR_EL0
+; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d11, d10, [sp, #32] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr w0, [x8, x0]
+; CHECK-NEXT:    ldp d13, d12, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr x30, [sp, #64] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp d15, d14, [sp], #80 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @x)
+  %1 = load i32, ptr %0, align 4
+  ret i32 %1
+}
+
+define i32 @load_tls_shared_za() nounwind "aarch64_inout_za" {
+; CHECK-LABEL: load_tls_shared_za:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp x29, x30, [sp, #-16]! // 16-byte Folded Spill
+; CHECK-NEXT:    mov x29, sp
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    rdsvl x8, #1
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    msub x9, x8, x8, x9
+; CHECK-NEXT:    mov sp, x9
+; CHECK-NEXT:    sub x10, x29, #16
+; CHECK-NEXT:    stp x9, x8, [x29, #-16]
+; CHECK-NEXT:    msr TPIDR2_EL0, x10
+; CHECK-NEXT:    adrp x0, :tlsdesc:x
+; CHECK-NEXT:    ldr x1, [x0, :tlsdesc_lo12:x]
+; CHECK-NEXT:    add x0, x0, :tlsdesc_lo12:x
+; CHECK-NEXT:    .tlsdesccall x
+; CHECK-NEXT:    blr x1
+; CHECK-NEXT:    mrs x8, TPIDR_EL0
+; CHECK-NEXT:    ldr w0, [x8, x0]
+; CHECK-NEXT:    mov w8, w0
+; CHECK-NEXT:    smstart za
+; CHECK-NEXT:    mrs x9, TPIDR2_EL0
+; CHECK-NEXT:    sub x0, x29, #16
+; CHECK-NEXT:    cbnz x9, .LBB2_2
+; CHECK-NEXT:  // %bb.1: // %entry
+; CHECK-NEXT:    bl __arm_tpidr2_restore
+; CHECK-NEXT:  .LBB2_2: // %entry
+; CHECK-NEXT:    mov w0, w8
+; CHECK-NEXT:    msr TPIDR2_EL0, xzr
+; CHECK-NEXT:    mov sp, x29
+; CHECK-NEXT:    ldp x29, x30, [sp], #16 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @x)
+  %1 = load i32, ptr %0, align 4
+  ret i32 %1
+}
+
+define i32 @load_tls_streaming_shared_za() nounwind "aarch64_inout_za" "aarch64_pstate_sm_enabled" {
+; CHECK-LABEL: load_tls_streaming_shared_za:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp d15, d14, [sp, #-96]! // 16-byte Folded Spill
+; CHECK-NEXT:    stp d13, d12, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d11, d10, [sp, #32] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d9, d8, [sp, #48] // 16-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #64] // 16-byte Folded Spill
+; CHECK-NEXT:    add x29, sp, #64
+; CHECK-NEXT:    str x19, [sp, #80] // 8-byte Folded Spill
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    rdsvl x8, #1
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    msub x9, x8, x8, x9
+; CHECK-NEXT:    mov sp, x9
+; CHECK-NEXT:    stp x9, x8, [x29, #-80]
+; CHECK-NEXT:    smstop sm
+; CHECK-NEXT:    sub x8, x29, #80
+; CHECK-NEXT:    msr TPIDR2_EL0, x8
+; CHECK-NEXT:    adrp x0, :tlsdesc:x
+; CHECK-NEXT:    ldr x1, [x0, :tlsdesc_lo12:x]
+; CHECK-NEXT:    add x0, x0, :tlsdesc_lo12:x
+; CHECK-NEXT:    .tlsdesccall x
+; CHECK-NEXT:    blr x1
+; CHECK-NEXT:    smstart sm
+; CHECK-NEXT:    mrs x8, TPIDR_EL0
+; CHECK-NEXT:    ldr w0, [x8, x0]
+; CHECK-NEXT:    mov w8, w0
+; CHECK-NEXT:    smstart za
+; CHECK-NEXT:    mrs x9, TPIDR2_EL0
+; CHECK-NEXT:    sub x0, x29, #80
+; CHECK-NEXT:    cbnz x9, .LBB3_2
+; CHECK-NEXT:  // %bb.1: // %entry
+; CHECK-NEXT:    bl __arm_tpidr2_restore
+; CHECK-NEXT:  .LBB3_2: // %entry
+; CHECK-NEXT:    mov w0, w8
+; CHECK-NEXT:    msr TPIDR2_EL0, xzr
+; CHECK-NEXT:    sub sp, x29, #64
+; CHECK-NEXT:    ldp x29, x30, [sp, #64] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr x19, [sp, #80] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d11, d10, [sp, #32] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d13, d12, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d15, d14, [sp], #96 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @x)
+  %1 = load i32, ptr %0, align 4
+  ret i32 %1
+}