[AArch64][SME] Disable tail-call optimization when streaming mode change or lazy-save may be required.

[sdesmalen]

Summary

When a streaming mode change is (or may be) required for a call, it will
need to restore the original mode after the call, which prevents the use of
tail-call optimization. The same holds true for a call that requires the lazy-save
mechanism to be set up before the call, and possibly restored after.

More details about the SME attributes and design can be found
in D131562.

[aemerson]

llvm/test/CodeGen/AArch64/sme-shared-za-interface.ll

8–10

This is correct as long as you have the shared-za attribute right?

[sdesmalen]

llvm/test/CodeGen/AArch64/sme-shared-za-interface.ll

8–10

Sorry, I missed this comment. The code below is only correct if @private_za_callee also has aarch64_pstate_za_shared, but in this case it isn't, which means that the call to @private_za_callee could clobber ZA. It would need the lazy-save mechanism (implemented in D133900)

[aemerson]

llvm/test/CodeGen/AArch64/sme-shared-za-interface.ll

8–10

Right, of course.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::isEligibleForTailCallOptimization(
   const SmallVector<ISD::OutputArg, 32> &Outs = CLI.Outs;
   const SmallVector<SDValue, 32> &OutVals = CLI.OutVals;
   const SmallVector<ISD::InputArg, 32> &Ins = CLI.Ins;
   const SelectionDAG &DAG = CLI.DAG;
   MachineFunction &MF = DAG.getMachineFunction();
   const Function &CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF.getCallingConv();
 
+  // SME Streaming functions are not eligible for TCO as they may require
+  // the streaming mode or ZA to be restored after returning from the call.
+  auto CallerAttrs = SMEAttrs::getFromFunction(MF.getFunction());
+  auto CalleeAttrs =
+      CLI.CB ? SMEAttrs::getFromCallBase(*CLI.CB) : SMEAttrs::getNormal();
+  if (CallerAttrs.requiresSMChange(CalleeAttrs) ||
+      CallerAttrs.requiresLazySave(CalleeAttrs))
+    return false;
+
   // Functions using the C or Fast calling convention that have an SVE signature
   // preserve more registers and should assume the SVE_VectorCall CC.
   // The check for matching callee-saved regs will determine whether it is
   // eligible for TCO.
   if ((CallerCC == CallingConv::C || CallerCC == CallingConv::Fast) &&
       MF.getInfo<AArch64FunctionInfo>()->isSVECC())
     CallerCC = CallingConv::AArch64_SVE_VectorCall;
 

llvm/test/CodeGen/AArch64/sme-shared-za-interface.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme  -verify-machineinstrs < %s | FileCheck %s
+
+declare void @private_za_callee()
+
+; Ensure that we don't use tail call optimization when a lazy-save is required.
+;
+; FIXME: The code below if obviously not yet correct, because it should set up
+; a lazy-save buffer before doing the call, and (conditionally) restore it after
+; the call. But this functionality will follow in a future patch.
+define void @disable_tailcallopt() "aarch64_pstate_za_shared" nounwind {
+; CHECK-LABEL: disable_tailcallopt:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    bl private_za_callee
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  tail call void @private_za_callee()
+  ret void
+}

llvm/test/CodeGen/AArch64/sme-streaming-compatible-interface.ll

 if.then:
   call void @streaming_callee()
   br label %exit
 
 exit:
   ret void
 }
 
+define void @disable_tailcallopt() "aarch64_pstate_sm_compatible" nounwind {
+; CHECK-LABEL: disable_tailcallopt:
+; CHECK:       // %bb.0:
+; 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:    stp x30, x19, [sp, #64] // 16-byte Folded Spill
+; CHECK-NEXT:    bl __arm_sme_state
+; CHECK-NEXT:    and x19, x0, #0x1
+; CHECK-NEXT:    tbz x19, #0, .LBB9_2
+; CHECK-NEXT:  // %bb.1:
+; CHECK-NEXT:    smstop sm
+; CHECK-NEXT:  .LBB9_2:
+; CHECK-NEXT:    bl normal_callee
+; CHECK-NEXT:    tbz x19, #0, .LBB9_4
+; CHECK-NEXT:  // %bb.3:
+; CHECK-NEXT:    smstart sm
+; CHECK-NEXT:  .LBB9_4:
+; CHECK-NEXT:    ldp x30, x19, [sp, #64] // 16-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], #80 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+
+  tail call void @normal_callee();
+  ret void;
+}
+
+
 attributes #0 = { nounwind "target-features"="+sve" }

llvm/test/CodeGen/AArch64/sme-streaming-interface.ll

 entry:
   %res = call fast double @llvm.cos.f64(double %x)
   %res.fadd = fadd fast double %res, %x
   ret double %res.fadd
 }
 
 declare double @llvm.cos.f64(double)
 
+; Ensure that tail call optimization is disabled when the streaming mode
+; doesn't match.
+define void @disable_tailcallopt() nounwind {
+; CHECK-LABEL: disable_tailcallopt:
+; CHECK:       // %bb.0:
+; 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:    smstart sm
+; CHECK-NEXT:    bl streaming_callee
+; CHECK-NEXT:    smstop sm
+; 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:    ldr x30, [sp, #64] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp d15, d14, [sp], #80 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+  tail call void @streaming_callee()
+  ret void;
+}
+
 attributes #0 = { nounwind "target-features"="+sve" }