[AArch64][SME] Disallow SME attributes on direct function calls

This was only used in a handful of tests (mainly to avoid making
multiple function declarations). These tests can easily be updated to
use indirect calls or attributes on declarations.

This allows us to remove checks that looked at both the "callee" and
"callsite" attributes, which makes the API of SMECallAttrs a clearer
and less error-prone (as you can't accidentally use .callee() when you
should have used .calleeOrCallsite()).

Note: This currently still allows non-conflicting attributes on direct
calls (as clang currently duplicates streaming mode attributes at each
callsite).

Author: MacDue

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8979,9 +8979,9 @@ static unsigned getSMCondition(const SMECallAttrs &CallAttrs) {
   if (!CallAttrs.caller().hasStreamingCompatibleInterface() ||
       CallAttrs.caller().hasStreamingBody())
     return AArch64SME::Always;
-  if (CallAttrs.calleeOrCallsite().hasNonStreamingInterface())
+  if (CallAttrs.callee().hasNonStreamingInterface())
     return AArch64SME::IfCallerIsStreaming;
-  if (CallAttrs.calleeOrCallsite().hasStreamingInterface())
+  if (CallAttrs.callee().hasStreamingInterface())
     return AArch64SME::IfCallerIsNonStreaming;
 
   llvm_unreachable("Unsupported attributes");
@@ -9158,7 +9158,7 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
       return DescribeCallsite(R) << " sets up a lazy save for ZA";
     });
   } else if (RequiresSaveAllZA) {
-    assert(!CallAttrs.calleeOrCallsite().hasSharedZAInterface() &&
+    assert(!CallAttrs.callee().hasSharedZAInterface() &&
            "Cannot share state that may not exist");
     Chain = emitSMEStateSaveRestore(*this, DAG, FuncInfo, DL, Chain,
                                     /*IsSave=*/true);
@@ -9484,9 +9484,9 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
       InGlue = Chain.getValue(1);
     }
 
-    SDValue NewChain = changeStreamingMode(
-        DAG, DL, CallAttrs.calleeOrCallsite().hasStreamingInterface(), Chain,
-        InGlue, getSMCondition(CallAttrs), PStateSM);
+    SDValue NewChain =
+        changeStreamingMode(DAG, DL, CallAttrs.callee().hasStreamingInterface(),
+                            Chain, InGlue, getSMCondition(CallAttrs), PStateSM);
     Chain = NewChain.getValue(0);
     InGlue = NewChain.getValue(1);
   }
@@ -9665,8 +9665,8 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   if (RequiresSMChange) {
     assert(PStateSM && "Expected a PStateSM to be set");
     Result = changeStreamingMode(
-        DAG, DL, !CallAttrs.calleeOrCallsite().hasStreamingInterface(), Result,
-        InGlue, getSMCondition(CallAttrs), PStateSM);
+        DAG, DL, !CallAttrs.callee().hasStreamingInterface(), Result, InGlue,
+        getSMCondition(CallAttrs), PStateSM);
 
     if (!Subtarget->isTargetDarwin() || Subtarget->hasSVE()) {
       InGlue = Result.getValue(1);

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -352,7 +352,7 @@ AArch64TTIImpl::getInlineCallPenalty(const Function *F, const CallBase &Call,
   SMEAttrs FAttrs(*F);
   SMECallAttrs CallAttrs(Call);
 
-  if (SMECallAttrs(FAttrs, CallAttrs.calleeOrCallsite()).requiresSMChange()) {
+  if (SMECallAttrs(FAttrs, CallAttrs.callee()).requiresSMChange()) {
     if (F == Call.getCaller()) // (1)
       return CallPenaltyChangeSM * DefaultCallPenalty;
     if (SMECallAttrs(FAttrs, CallAttrs.caller()).requiresSMChange()) // (2)

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

@@ -94,22 +94,28 @@ void SMEAttrs::addKnownFunctionAttrs(StringRef FuncName) {
 }
 
 bool SMECallAttrs::requiresSMChange() const {
-  if ((Callsite | Callee).hasStreamingCompatibleInterface())
+  if (callee().hasStreamingCompatibleInterface())
     return false;
 
   // Both non-streaming
-  if (Caller.hasNonStreamingInterfaceAndBody() &&
-      (Callsite | Callee).hasNonStreamingInterface())
+  if (caller().hasNonStreamingInterfaceAndBody() &&
+      callee().hasNonStreamingInterface())
     return false;
 
   // Both streaming
-  if (Caller.hasStreamingInterfaceOrBody() &&
-      (Callsite | Callee).hasStreamingInterface())
+  if (caller().hasStreamingInterfaceOrBody() &&
+      callee().hasStreamingInterface())
     return false;
 
   return true;
 }
 
 SMECallAttrs::SMECallAttrs(const CallBase &CB)
-    : SMECallAttrs(*CB.getFunction(), CB.getCalledFunction(),
-                   CB.getAttributes()) {}
+    : CallerFn(*CB.getFunction()), CalledFn(CB.getCalledFunction()),
+      Callsite(CB.getAttributes()), IsIndirect(CB.isIndirectCall()) {
+  // FIXME: We probably should not allow SME attributes on direct calls but
+  // clang currently copies attributes from the declaration to each callsite.
+  assert((IsIndirect || Callsite.withoutPerCallsiteFlags().isNormal() ||
+          Callsite.withoutPerCallsiteFlags() == CalledFn) &&
+         "SME attributes at callsite do not match declaration");
+}

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

@@ -44,7 +44,8 @@ class SMEAttrs {
     ZA_Shift = 6,
     ZA_Mask = 0b111 << ZA_Shift,
     ZT0_Shift = 9,
-    ZT0_Mask = 0b111 << ZT0_Shift
+    ZT0_Mask = 0b111 << ZT0_Shift,
+    Callsite_Flags = ZT0_Undef
   };
 
   SMEAttrs() = default;
@@ -129,10 +130,13 @@ class SMEAttrs {
   }
   bool hasZT0State() const { return isNewZT0() || sharesZT0(); }
 
-  SMEAttrs operator|(SMEAttrs Other) const {
-    SMEAttrs Merged(*this);
-    Merged.set(Other.Bitmask, /*Enable=*/true);
-    return Merged;
+  bool isNormal() const { return Bitmask == Normal; }
+  SMEAttrs withoutPerCallsiteFlags() const {
+    return (Bitmask & ~Callsite_Flags);
+  }
+
+  bool operator==(SMEAttrs const &Other) const {
+    return Bitmask == Other.Bitmask;
   }
 
 private:
@@ -143,54 +147,57 @@ class SMEAttrs {
 /// interfaces to query whether a streaming mode change or lazy-save mechanism
 /// is required when going from one function to another (e.g. through a call).
 class SMECallAttrs {
-  SMEAttrs Caller;
-  SMEAttrs Callee;
+  SMEAttrs CallerFn;
+  SMEAttrs CalledFn;
   SMEAttrs Callsite;
+  bool IsIndirect = false;
 
 public:
   SMECallAttrs(SMEAttrs Caller, SMEAttrs Callee,
                SMEAttrs Callsite = SMEAttrs::Normal)
-      : Caller(Caller), Callee(Callee), Callsite(Callsite) {}
+      : CallerFn(Caller), CalledFn(Callee), Callsite(Callsite) {}
 
   SMECallAttrs(const CallBase &CB);
 
-  SMEAttrs &caller() { return Caller; }
-  SMEAttrs &callee() { return Callee; }
+  SMEAttrs &caller() { return CallerFn; }
+  SMEAttrs &callee() {
+    if (IsIndirect)
+      return Callsite;
+    return CalledFn;
+  }
   SMEAttrs &callsite() { return Callsite; }
-  SMEAttrs const &caller() const { return Caller; }
-  SMEAttrs const &callee() const { return Callee; }
+  SMEAttrs const &caller() const { return CallerFn; }
+  SMEAttrs const &callee() const {
+    return const_cast<SMECallAttrs *>(this)->callee();
+  }
   SMEAttrs const &callsite() const { return Callsite; }
-  SMEAttrs calleeOrCallsite() const { return Callsite | Callee; }
 
   /// \return true if a call from Caller -> Callee requires a change in
   /// streaming mode.
   bool requiresSMChange() const;
 
   bool requiresLazySave() const {
-    return Caller.hasZAState() && (Callsite | Callee).hasPrivateZAInterface() &&
-           !Callee.isSMEABIRoutine();
+    return caller().hasZAState() && callee().hasPrivateZAInterface() &&
+           !callee().isSMEABIRoutine();
   }
 
   bool requiresPreservingZT0() const {
-    return Caller.hasZT0State() && !Callsite.hasUndefZT0() &&
-           !(Callsite | Callee).sharesZT0() &&
-           !(Callsite | Callee).hasAgnosticZAInterface();
+    return caller().hasZT0State() && !callsite().hasUndefZT0() &&
+           !callee().sharesZT0() && !callee().hasAgnosticZAInterface();
   }
 
   bool requiresDisablingZABeforeCall() const {
-    return Caller.hasZT0State() && !Caller.hasZAState() &&
-           (Callsite | Callee).hasPrivateZAInterface() &&
-           !Callee.isSMEABIRoutine();
+    return caller().hasZT0State() && !caller().hasZAState() &&
+           callee().hasPrivateZAInterface() && !callee().isSMEABIRoutine();
   }
 
   bool requiresEnablingZAAfterCall() const {
     return requiresLazySave() || requiresDisablingZABeforeCall();
   }
 
   bool requiresPreservingAllZAState() const {
-    return Caller.hasAgnosticZAInterface() &&
-           !(Callsite | Callee).hasAgnosticZAInterface() &&
-           !Callee.isSMEABIRoutine();
+    return caller().hasAgnosticZAInterface() &&
+           !callee().hasAgnosticZAInterface() && !callee().isSMEABIRoutine();
   }
 };
 

llvm/test/CodeGen/AArch64/sme-peephole-opts.ll

@@ -2,11 +2,12 @@
 ; RUN: llc -mtriple=aarch64-linux-gnu -aarch64-streaming-hazard-size=0 -mattr=+sve,+sme2 < %s | FileCheck %s
 
 declare void @callee()
+declare void @callee_sm() "aarch64_pstate_sm_enabled"
 declare void @callee_farg(float)
 declare float @callee_farg_fret(float)
 
 ; normal caller -> streaming callees
-define void @test0() nounwind {
+define void @test0(ptr %callee) nounwind {
 ; CHECK-LABEL: test0:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    stp d15, d14, [sp, #-80]! // 16-byte Folded Spill
@@ -16,17 +17,17 @@ define void @test0() nounwind {
 ; CHECK-NEXT:    stp d9, d8, [sp, #48] // 16-byte Folded Spill
 ; CHECK-NEXT:    stp x30, x9, [sp, #64] // 16-byte Folded Spill
 ; CHECK-NEXT:    smstart sm
-; CHECK-NEXT:    bl callee
-; CHECK-NEXT:    bl callee
+; CHECK-NEXT:    bl callee_sm
+; CHECK-NEXT:    bl callee_sm
 ; CHECK-NEXT:    smstop sm
 ; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr x30, [sp, #64] // 8-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], #80 // 16-byte Folded Reload
 ; CHECK-NEXT:    ret
-  call void @callee() "aarch64_pstate_sm_enabled"
-  call void @callee() "aarch64_pstate_sm_enabled"
+  call void @callee_sm()
+  call void @callee_sm()
   ret void
 }
 
@@ -118,7 +119,7 @@ define void @test3() nounwind "aarch64_pstate_sm_compatible" {
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    smstart sm
 ; CHECK-NEXT:  .LBB3_2:
-; CHECK-NEXT:    bl callee
+; CHECK-NEXT:    bl callee_sm
 ; CHECK-NEXT:    tbnz w19, #0, .LBB3_4
 ; CHECK-NEXT:  // %bb.3:
 ; CHECK-NEXT:    smstop sm
@@ -140,7 +141,7 @@ define void @test3() nounwind "aarch64_pstate_sm_compatible" {
 ; CHECK-NEXT:  // %bb.9:
 ; CHECK-NEXT:    smstart sm
 ; CHECK-NEXT:  .LBB3_10:
-; CHECK-NEXT:    bl callee
+; CHECK-NEXT:    bl callee_sm
 ; CHECK-NEXT:    tbnz w19, #0, .LBB3_12
 ; CHECK-NEXT:  // %bb.11:
 ; CHECK-NEXT:    smstop sm
@@ -152,9 +153,9 @@ define void @test3() nounwind "aarch64_pstate_sm_compatible" {
 ; 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
-  call void @callee() "aarch64_pstate_sm_enabled"
+  call void @callee_sm()
   call void @callee()
-  call void @callee() "aarch64_pstate_sm_enabled"
+  call void @callee_sm()
   ret void
 }
 
@@ -342,7 +343,7 @@ define void @test10() "aarch64_pstate_sm_body" {
 ; CHECK-NEXT:    bl callee
 ; CHECK-NEXT:    smstart sm
 ; CHECK-NEXT:    .cfi_restore vg
-; CHECK-NEXT:    bl callee
+; CHECK-NEXT:    bl callee_sm
 ; CHECK-NEXT:    .cfi_offset vg, -24
 ; CHECK-NEXT:    smstop sm
 ; CHECK-NEXT:    bl callee
@@ -363,7 +364,7 @@ define void @test10() "aarch64_pstate_sm_body" {
 ; CHECK-NEXT:    .cfi_restore b15
 ; CHECK-NEXT:    ret
   call void @callee()
-  call void @callee() "aarch64_pstate_sm_enabled"
+  call void @callee_sm()
   call void @callee()
   ret void
 }

llvm/test/CodeGen/AArch64/sme-vg-to-stack.ll

@@ -1098,11 +1098,11 @@ define void @test_rdsvl_right_after_prologue(i64 %x0) nounwind {
 ; NO-SVE-CHECK-NEXT: ret
   %some_alloc = alloca i64, align 8
   %rdsvl = tail call i64 @llvm.aarch64.sme.cntsd()
-  call void @bar(i64 %rdsvl, i64 %x0) "aarch64_pstate_sm_enabled"
+  call void @bar(i64 %rdsvl, i64 %x0)
   ret void
 }
 
-declare void @bar(i64, i64)
+declare void @bar(i64, i64) "aarch64_pstate_sm_enabled"
 
 ; Ensure we still emit async unwind information with -fno-asynchronous-unwind-tables
 ; if the function contains a streaming-mode change.