[clang][SME] Ignore flatten/clang::always_inline statements for callees with mismatched streaming attributes

clang/lib/CodeGen/CGCall.cpp

@@ -5112,9 +5112,10 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
 
   // Some architectures (such as x86-64) have the ABI changed based on
   // attribute-target/features. Give them a chance to diagnose.
-  CGM.getTargetCodeGenInfo().checkFunctionCallABI(
-      CGM, Loc, dyn_cast_or_null<FunctionDecl>(CurCodeDecl),
-      dyn_cast_or_null<FunctionDecl>(TargetDecl), CallArgs, RetTy);
+  const FunctionDecl *CallerDecl = dyn_cast_or_null<FunctionDecl>(CurCodeDecl);
+  const FunctionDecl *CalleeDecl = dyn_cast_or_null<FunctionDecl>(TargetDecl);
+  CGM.getTargetCodeGenInfo().checkFunctionCallABI(CGM, Loc, CallerDecl,
+                                                  CalleeDecl, CallArgs, RetTy);
 
   // 1. Set up the arguments.
 
@@ -5705,7 +5706,9 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   // FIXME: should this really take priority over __try, below?
   if (CurCodeDecl && CurCodeDecl->hasAttr<FlattenAttr>() &&
       !InNoInlineAttributedStmt &&
-      !(TargetDecl && TargetDecl->hasAttr<NoInlineAttr>())) {
+      !(TargetDecl && TargetDecl->hasAttr<NoInlineAttr>()) &&
+      !CGM.getTargetCodeGenInfo().wouldInliningViolateFunctionCallABI(
+          CallerDecl, CalleeDecl)) {
     Attrs =
         Attrs.addFnAttribute(getLLVMContext(), llvm::Attribute::AlwaysInline);
   }

clang/lib/CodeGen/TargetInfo.h

@@ -98,6 +98,15 @@ class TargetCodeGenInfo {
                                     const CallArgList &Args,
                                     QualType ReturnType) const {}
 
+  /// Returns true if inlining the function call would produce incorrect code
+  /// for the current target and should be ignored (even with the always_inline
+  /// or flatten attributes).
+  virtual bool
+  wouldInliningViolateFunctionCallABI(const FunctionDecl *Caller,
+                                      const FunctionDecl *Callee) const {
+    return false;
+  }
+
   /// Determines the size of struct _Unwind_Exception on this platform,
   /// in 8-bit units.  The Itanium ABI defines this as:
   ///   struct _Unwind_Exception {

clang/lib/CodeGen/Targets/AArch64.cpp

@@ -177,6 +177,9 @@ class AArch64TargetCodeGenInfo : public TargetCodeGenInfo {
                             const FunctionDecl *Callee, const CallArgList &Args,
                             QualType ReturnType) const override;
 
+  bool wouldInliningViolateFunctionCallABI(
+      const FunctionDecl *Caller, const FunctionDecl *Callee) const override;
+
 private:
   // Diagnose calls between functions with incompatible Streaming SVE
   // attributes.
@@ -1143,30 +1146,62 @@ void AArch64TargetCodeGenInfo::checkFunctionABI(
   }
 }
 
-void AArch64TargetCodeGenInfo::checkFunctionCallABIStreaming(
-    CodeGenModule &CGM, SourceLocation CallLoc, const FunctionDecl *Caller,
-    const FunctionDecl *Callee) const {
-  if (!Caller || !Callee || !Callee->hasAttr<AlwaysInlineAttr>())
-    return;
+enum class ArmStreamingInlinability : uint8_t {
+  Ok = 0,
+  IncompatibleStreamingModes = 1,
+  MismatchedStreamingCompatibility = 1 << 1,
+  CalleeHasNewZA = 1 << 2,
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/CalleeHasNewZA),
+};
 
+/// Determines if there are any streaming ABI issues with inlining \p Callee
+/// into \p Caller. Returns the issues in the ArmStreamingInlinability bit enum
+/// (multiple bits can be set).
+static ArmStreamingInlinability
+GetArmStreamingInlinability(const FunctionDecl *Caller,
+                            const FunctionDecl *Callee) {
   bool CallerIsStreaming =
       IsArmStreamingFunction(Caller, /*IncludeLocallyStreaming=*/true);
   bool CalleeIsStreaming =
       IsArmStreamingFunction(Callee, /*IncludeLocallyStreaming=*/true);
   bool CallerIsStreamingCompatible = isStreamingCompatible(Caller);
   bool CalleeIsStreamingCompatible = isStreamingCompatible(Callee);
 
+  ArmStreamingInlinability Inlinability = ArmStreamingInlinability::Ok;
+
   if (!CalleeIsStreamingCompatible &&
-      (CallerIsStreaming != CalleeIsStreaming || CallerIsStreamingCompatible))
+      (CallerIsStreaming != CalleeIsStreaming || CallerIsStreamingCompatible)) {
+    Inlinability |= ArmStreamingInlinability::MismatchedStreamingCompatibility;
+    if (CalleeIsStreaming)
+      Inlinability |= ArmStreamingInlinability::IncompatibleStreamingModes;
+  }
+  if (auto *NewAttr = Callee->getAttr<ArmNewAttr>())
+    if (NewAttr->isNewZA())
+      Inlinability |= ArmStreamingInlinability::CalleeHasNewZA;
+
+  return Inlinability;
+}
+
+void AArch64TargetCodeGenInfo::checkFunctionCallABIStreaming(
+    CodeGenModule &CGM, SourceLocation CallLoc, const FunctionDecl *Caller,
+    const FunctionDecl *Callee) const {
+  if (!Caller || !Callee || !Callee->hasAttr<AlwaysInlineAttr>())
+    return;
+
+  ArmStreamingInlinability Inlinability =
+      GetArmStreamingInlinability(Caller, Callee);
+
+  if (bool(Inlinability &
+           ArmStreamingInlinability::MismatchedStreamingCompatibility))
     CGM.getDiags().Report(
-        CallLoc, CalleeIsStreaming
+        CallLoc, bool(Inlinability &
+                      ArmStreamingInlinability::IncompatibleStreamingModes)
                      ? diag::err_function_always_inline_attribute_mismatch
                      : diag::warn_function_always_inline_attribute_mismatch)
         << Caller->getDeclName() << Callee->getDeclName() << "streaming";
-  if (auto *NewAttr = Callee->getAttr<ArmNewAttr>())
-    if (NewAttr->isNewZA())
-      CGM.getDiags().Report(CallLoc, diag::err_function_always_inline_new_za)
-          << Callee->getDeclName();
+  if (bool(Inlinability & ArmStreamingInlinability::CalleeHasNewZA))
+    CGM.getDiags().Report(CallLoc, diag::err_function_always_inline_new_za)
+        << Callee->getDeclName();
 }
 
 // If the target does not have floating-point registers, but we are using a
@@ -1200,6 +1235,13 @@ void AArch64TargetCodeGenInfo::checkFunctionCallABI(CodeGenModule &CGM,
   checkFunctionCallABISoftFloat(CGM, CallLoc, Caller, Callee, Args, ReturnType);
 }
 
+bool AArch64TargetCodeGenInfo::wouldInliningViolateFunctionCallABI(
+    const FunctionDecl *Caller, const FunctionDecl *Callee) const {
+  return Caller && Callee &&
+         GetArmStreamingInlinability(Caller, Callee) !=
+             ArmStreamingInlinability::Ok;
+}
+
 void AArch64ABIInfo::appendAttributeMangling(TargetClonesAttr *Attr,
                                              unsigned Index,
                                              raw_ostream &Out) const {

clang/test/CodeGen/AArch64/sme-flatten-streaming-attrs.c

@@ -0,0 +1,74 @@
+// RUN: %clang_cc1 -triple aarch64-none-linux-gnu -emit-llvm -target-feature +sme %s -o - | FileCheck %s
+
+// REQUIRES: aarch64-registered-target
+
+extern void was_inlined(void);
+
+#define __flatten  __attribute__((flatten))
+void fn(void) { was_inlined(); }
+void fn_streaming_compatible(void) __arm_streaming_compatible { was_inlined(); }
+void fn_streaming(void) __arm_streaming { was_inlined(); }
+__arm_locally_streaming void fn_locally_streaming(void) { was_inlined(); }
+__arm_new("za") void fn_streaming_new_za(void) __arm_streaming { was_inlined(); }
+
+__flatten
+void caller(void) {
+    fn();
+    fn_streaming_compatible();
+    fn_streaming();
+    fn_locally_streaming();
+    fn_streaming_new_za();
+}
+// CHECK-LABEL: void @caller()
+//  CHECK-NEXT: entry:
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @fn_streaming
+//  CHECK-NEXT:   call void @fn_locally_streaming
+//  CHECK-NEXT:   call void @fn_streaming_new_za
+
+__flatten void caller_streaming_compatible(void) __arm_streaming_compatible {
+    fn();
+    fn_streaming_compatible();
+    fn_streaming();
+    fn_locally_streaming();
+    fn_streaming_new_za();
+}
+// CHECK-LABEL: void @caller_streaming_compatible()
+//  CHECK-NEXT: entry:
+//  CHECK-NEXT:   call void @fn
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @fn_streaming
+//  CHECK-NEXT:   call void @fn_locally_streaming
+//  CHECK-NEXT:   call void @fn_streaming_new_za
+
+__flatten void caller_streaming(void) __arm_streaming {
+    fn();
+    fn_streaming_compatible();
+    fn_streaming();
+    fn_locally_streaming();
+    fn_streaming_new_za();
+}
+// CHECK-LABEL: void @caller_streaming()
+//  CHECK-NEXT: entry:
+//  CHECK-NEXT:   call void @fn
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @fn_streaming_new_za
+
+__flatten __arm_locally_streaming
+void caller_locally_streaming(void) {
+    fn();
+    fn_streaming_compatible();
+    fn_streaming();
+    fn_locally_streaming();
+    fn_streaming_new_za();
+}
+// CHECK-LABEL: void @caller_locally_streaming()
+//  CHECK-NEXT: entry:
+//  CHECK-NEXT:   call void @fn
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @was_inlined
+//  CHECK-NEXT:   call void @fn_streaming_new_za