[clang][SME] Account for C++ lambdas in SME builtin diagnostics

clang/lib/Sema/SemaARM.cpp

@@ -650,7 +650,8 @@ static ArmSMEState getSMEState(unsigned BuiltinID) {
 
 bool SemaARM::CheckSMEBuiltinFunctionCall(unsigned BuiltinID,
                                           CallExpr *TheCall) {
-  if (const FunctionDecl *FD = SemaRef.getCurFunctionDecl()) {
+  if (const FunctionDecl *FD =
+          SemaRef.getCurFunctionDecl(/*AllowLambda=*/true)) {
     std::optional<ArmStreamingType> BuiltinType;
 
     switch (BuiltinID) {
@@ -690,7 +691,8 @@ bool SemaARM::CheckSMEBuiltinFunctionCall(unsigned BuiltinID,
 
 bool SemaARM::CheckSVEBuiltinFunctionCall(unsigned BuiltinID,
                                           CallExpr *TheCall) {
-  if (const FunctionDecl *FD = SemaRef.getCurFunctionDecl()) {
+  if (const FunctionDecl *FD =
+          SemaRef.getCurFunctionDecl(/*AllowLambda=*/true)) {
     std::optional<ArmStreamingType> BuiltinType;
 
     switch (BuiltinID) {
@@ -719,7 +721,8 @@ bool SemaARM::CheckSVEBuiltinFunctionCall(unsigned BuiltinID,
 bool SemaARM::CheckNeonBuiltinFunctionCall(const TargetInfo &TI,
                                            unsigned BuiltinID,
                                            CallExpr *TheCall) {
-  if (const FunctionDecl *FD = SemaRef.getCurFunctionDecl()) {
+  if (const FunctionDecl *FD =
+          SemaRef.getCurFunctionDecl(/*AllowLambda=*/true)) {
 
     switch (BuiltinID) {
     default:

clang/test/Sema/aarch64-incompat-sm-builtin-calls.c → clang/test/Sema/aarch64-incompat-sm-builtin-calls.cpp

@@ -1,5 +1,5 @@
 // NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py
-// RUN: %clang_cc1  -triple aarch64-none-linux-gnu -target-feature +sve \
+// RUN: %clang_cc1  -std=c++23 -triple aarch64-none-linux-gnu -target-feature +sve \
 // RUN:   -target-feature +bf16 -target-feature +sve -target-feature +sme -target-feature +sme2 -target-feature +sve2 -target-feature +neon -Waarch64-sme-attributes -fsyntax-only -verify %s
 
 // REQUIRES: aarch64-registered-target
@@ -126,3 +126,69 @@ void missing_zt0(void) __arm_streaming {
 
 __arm_new("zt0")
 void new_zt0(void) __arm_streaming { svzero_zt(0); }
+
+/// C++ lambda tests:
+
+void use_streaming_builtin_in_lambda(uint32_t slice_base, svbool_t pg, const void *ptr) __arm_streaming __arm_out("za")
+{
+  [&]{
+    /// The lambda is its own function and does not inherit the SME attributes (so this should error).
+    // expected-error@+1 {{builtin can only be called from a streaming function}}
+    svld1_hor_za64(0, slice_base, pg, ptr);
+  }();
+}
+
+void use_streaming_builtin(uint32_t slice_base, svbool_t pg, const void *ptr) __arm_streaming __arm_out("za")
+{
+  /// Without the lambda the same builtin is okay (as the SME attributes apply).
+  svld1_hor_za64(0, slice_base, pg, ptr);
+}
+
+int16x8_t use_neon_builtin_sm(int16x8_t splat) __arm_streaming_compatible {
+  // expected-error@+1 {{builtin can only be called from a non-streaming function}}
+  return (int16x8_t)__builtin_neon_vqaddq_v((int8x16_t)splat, (int8x16_t)splat, 33);
+}
+
+int16x8_t use_neon_builtin_sm_in_lambda(int16x8_t splat) __arm_streaming_compatible {
+  return [&]{
+    /// This should not error (as we switch out of streaming mode to execute the lambda).
+    /// Note: The result int16x8_t is spilled and reloaded as a q-register.
+    return (int16x8_t)__builtin_neon_vqaddq_v((int8x16_t)splat, (int8x16_t)splat, 33);
+  }();
+}
+
+float use_incomp_sve_builtin_sm() __arm_streaming {
+  // expected-error@+1 {{builtin can only be called from a non-streaming function}}
+  return svadda(svptrue_b32(), 0, svdup_f32(1));
+}
+
+float incomp_sve_sm_fadda_sm_in_lambda(void) __arm_streaming {
+  return [&]{
+    /// This should work like the Neon builtin.
+    return svadda(svptrue_b32(), 0, svdup_f32(1));
+  }();
+}
+
+void use_streaming_builtin_in_streaming_lambda(uint32_t slice_base, const void *ptr)
+{
+  [&]  __arm_new("za") () __arm_streaming {
+    // Here the lambda is streaming with ZA state, so this is okay.
+    svld1_hor_za64(0, slice_base, svptrue_b64(), ptr);
+  }();
+}
+
+int16x8_t use_neon_builtin_in_streaming_lambda(int16x8_t splat) {
+  return [&]() __arm_streaming_compatible {
+    /// This should error as the lambda is streaming-compatible.
+    // expected-error@+1 {{builtin can only be called from a non-streaming function}}
+    return (int16x8_t)__builtin_neon_vqaddq_v((int8x16_t)splat, (int8x16_t)splat, 33);
+  }();
+}
+
+float incomp_sve_fadda_in_streaming_lambda(void) {
+  return [&]() __arm_streaming {
+    // Should error (like the Neon case above).
+    // expected-error@+1 {{builtin can only be called from a non-streaming function}}
+    return svadda(svptrue_b32(), 0, svdup_f32(1));
+  }();
+}