[clang/LLVM] Add flatten_deep attribute for depth-limited inlining (1/2)

clang/include/clang/Basic/Attr.td

@@ -1984,6 +1984,13 @@ def Flatten : InheritableAttr {
   let SimpleHandler = 1;
 }
 
+def FlattenDeep : InheritableAttr {
+  let Spellings = [Clang<"flatten_deep">];
+  let Subjects = SubjectList<[Function], ErrorDiag>;
+  let Args = [UnsignedArgument<"MaxDepth">];
+  let Documentation = [FlattenDeepDocs];
+}
+
 def Format : InheritableAttr {
   let Spellings = [GCC<"format">];
   let Args = [IdentifierArgument<"Type">, IntArgument<"FormatIdx">,

clang/include/clang/Basic/AttrDocs.td

@@ -4032,6 +4032,29 @@ callee is unavailable or if the callee has the ``noinline`` attribute.
   }];
 }
 
+def FlattenDeepDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The ``flatten_deep`` attribute causes calls within the attributed function and
+their transitive callees to be inlined up to a specified depth, unless it is
+impossible to do so (for example if the body of the callee is unavailable or if
+the callee has the ``noinline`` attribute).
+
+This attribute takes a single unsigned integer argument representing the maximum
+depth of the call tree to inline. For example, ``__attribute__((flatten_deep(3)))``
+will inline all calls within the function, then inline all calls within those
+inlined functions (depth 2), and then inline all calls within those functions
+(depth 3).
+
+.. code-block:: c++
+
+  __attribute__((flatten_deep(3)))
+  void process_data() {
+    // All calls up to 3 levels deep in the call tree will be inlined
+  }
+  }];
+}
+
 def FormatDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2749,6 +2749,15 @@ void CodeGenModule::SetLLVMFunctionAttributesForDefinition(const Decl *D,
       B.addAttribute("aarch64_new_zt0");
   }
 
+  // Handle flatten_deep attribute for depth-based inlining
+  if (const auto *FD = dyn_cast<FunctionDecl>(D)) {
+    if (const FlattenDeepAttr *FDA = FD->getAttr<FlattenDeepAttr>()) {
+      // Add the flatten_deep attribute with the max depth value as a typed int
+      // attribute
+      B.addRawIntAttr(llvm::Attribute::FlattenDeep, FDA->getMaxDepth());
+    }
+  }
+
   // Track whether we need to add the optnone LLVM attribute,
   // starting with the default for this optimization level.
   bool ShouldAddOptNone =

clang/lib/Sema/SemaDeclAttr.cpp

@@ -3695,6 +3695,24 @@ static void handleInitPriorityAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   D->addAttr(::new (S.Context) InitPriorityAttr(S.Context, AL, prioritynum));
 }
 
+static void handleFlattenDeepAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
+  Expr *E = AL.getArgAsExpr(0);
+  uint32_t maxDepth;
+  if (!S.checkUInt32Argument(AL, E, maxDepth)) {
+    AL.setInvalid();
+    return;
+  }
+
+  if (maxDepth == 0) {
+    S.Diag(AL.getLoc(), diag::err_attribute_argument_is_zero)
+        << AL << E->getSourceRange();
+    AL.setInvalid();
+    return;
+  }
+
+  D->addAttr(::new (S.Context) FlattenDeepAttr(S.Context, AL, maxDepth));
+}
+
 ErrorAttr *Sema::mergeErrorAttr(Decl *D, const AttributeCommonInfo &CI,
                                 StringRef NewUserDiagnostic) {
   if (const auto *EA = D->getAttr<ErrorAttr>()) {
@@ -7236,6 +7254,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_Format:
     handleFormatAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_FlattenDeep:
+    handleFlattenDeepAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_FormatMatches:
     handleFormatMatchesAttr(S, D, AL);
     break;

clang/test/CodeGen/attr-flatten-deep-cg.c

@@ -0,0 +1,10 @@
+// RUN: %clang_cc1 -emit-llvm -o - %s | FileCheck %s
+
+// CHECK-LABEL: define {{.*}} @test1
+// CHECK-SAME: #[[ATTR1:[0-9]+]]
+__attribute__((flatten_deep(3)))
+void test1() {
+}
+
+// Verify the attribute is present in the attribute groups
+// CHECK-DAG: attributes #[[ATTR1]] = { {{.*}}flatten_deep=3{{.*}} }

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -86,6 +86,7 @@
 // CHECK-NEXT: ExternalSourceSymbol ((SubjectMatchRule_record, SubjectMatchRule_enum, SubjectMatchRule_enum_constant, SubjectMatchRule_field, SubjectMatchRule_function, SubjectMatchRule_namespace, SubjectMatchRule_objc_category, SubjectMatchRule_objc_implementation, SubjectMatchRule_objc_interface, SubjectMatchRule_objc_method, SubjectMatchRule_objc_property, SubjectMatchRule_objc_protocol, SubjectMatchRule_record, SubjectMatchRule_type_alias, SubjectMatchRule_variable))
 // CHECK-NEXT: FlagEnum (SubjectMatchRule_enum)
 // CHECK-NEXT: Flatten (SubjectMatchRule_function)
+// CHECK-NEXT: FlattenDeep (SubjectMatchRule_function)
 // CHECK-NEXT: FunctionReturnThunks (SubjectMatchRule_function)
 // CHECK-NEXT: GNUInline (SubjectMatchRule_function)
 // CHECK-NEXT: HIPManaged (SubjectMatchRule_variable)

clang/test/Sema/attr-flatten-deep.c

@@ -0,0 +1,14 @@
+// RUN: %clang_cc1 -fsyntax-only -verify %s
+
+// Test basic usage - valid
+__attribute__((flatten_deep(3)))
+void test_valid() {
+}
+
+// Test attribute on non-function - should error
+__attribute__((flatten_deep(3))) int x; // expected-error {{'flatten_deep' attribute only applies to functions}}
+
+// Test depth = 0 - should error (depth must be >= 1)
+__attribute__((flatten_deep(0))) // expected-error {{'flatten_deep' attribute must be greater than 0}}
+void test_depth_zero() {
+}

llvm/include/llvm/AsmParser/LLParser.h

@@ -320,6 +320,7 @@ namespace llvm {
                                 bool AllowParens = false);
     bool parseOptionalCodeModel(CodeModel::Model &model);
     bool parseOptionalDerefAttrBytes(lltok::Kind AttrKind, uint64_t &Bytes);
+    bool parseOptionalFlattenDeepDepth(lltok::Kind AttrKind, uint64_t &Depth);
     bool parseOptionalUWTableKind(UWTableKind &Kind);
     bool parseAllocKind(AllocFnKind &Kind);
     std::optional<MemoryEffects> parseMemoryAttr();

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -801,6 +801,7 @@ enum AttributeKindCodes {
   ATTR_KIND_CAPTURES = 102,
   ATTR_KIND_DEAD_ON_RETURN = 103,
   ATTR_KIND_SANITIZE_ALLOC_TOKEN = 104,
+  ATTR_KIND_FLATTEN_DEEP = 105,
 };
 
 enum ComdatSelectionKindCodes {

llvm/include/llvm/IR/Attributes.h

@@ -1246,6 +1246,10 @@ class AttrBuilder {
   /// form used internally in Attribute.
   LLVM_ABI AttrBuilder &addDereferenceableOrNullAttr(uint64_t Bytes);
 
+  /// This turns the flatten_deep depth into the form used internally in
+  /// Attribute.
+  LLVM_ABI AttrBuilder &addFlattenDeepAttr(uint64_t Depth);
+
   /// This turns one (or two) ints into the form used internally in Attribute.
   LLVM_ABI AttrBuilder &
   addAllocSizeAttr(unsigned ElemSizeArg,

llvm/include/llvm/IR/Attributes.td

@@ -207,6 +207,9 @@ def NoImplicitFloat : EnumAttr<"noimplicitfloat", IntersectPreserve, [FnAttr]>;
 /// inline=never.
 def NoInline : EnumAttr<"noinline", IntersectPreserve, [FnAttr]>;
 
+/// Inline calls recursively up to specified depth.
+def FlattenDeep : IntAttr<"flatten_deep", IntersectPreserve, [FnAttr]>;
+
 /// Function is called early and/or often, so lazy binding isn't worthwhile.
 def NonLazyBind : EnumAttr<"nonlazybind", IntersectPreserve, [FnAttr]>;
 

llvm/lib/AsmParser/LLParser.cpp

@@ -1620,6 +1620,22 @@ bool LLParser::parseEnumAttribute(Attribute::AttrKind Attr, AttrBuilder &B,
     B.addDereferenceableOrNullAttr(Bytes);
     return false;
   }
+  case Attribute::FlattenDeep: {
+    uint64_t Depth = 0;
+    if (InAttrGroup) {
+      Lex.Lex();
+      LocTy DepthLoc = Lex.getLoc();
+      if (parseToken(lltok::equal, "expected '=' here") || parseUInt64(Depth))
+        return true;
+      if (!Depth)
+        return error(DepthLoc, "flatten_deep depth must be non-zero");
+    } else {
+      if (parseOptionalFlattenDeepDepth(lltok::kw_flatten_deep, Depth))
+        return true;
+    }
+    B.addFlattenDeepAttr(Depth);
+    return false;
+  }
   case Attribute::UWTable: {
     UWTableKind Kind;
     if (parseOptionalUWTableKind(Kind))
@@ -2494,6 +2510,30 @@ bool LLParser::parseOptionalDerefAttrBytes(lltok::Kind AttrKind,
   return false;
 }
 
+/// parseOptionalFlattenDeepDepth
+///   ::= /* empty */
+///   ::= 'flatten_deep' '(' 4 ')'
+bool LLParser::parseOptionalFlattenDeepDepth(lltok::Kind AttrKind,
+                                             uint64_t &Depth) {
+  assert(AttrKind == lltok::kw_flatten_deep && "contract!");
+
+  Depth = 0;
+  if (!EatIfPresent(AttrKind))
+    return false;
+  LocTy ParenLoc = Lex.getLoc();
+  if (!EatIfPresent(lltok::lparen))
+    return error(ParenLoc, "expected '('");
+  LocTy DepthLoc = Lex.getLoc();
+  if (parseUInt64(Depth))
+    return true;
+  ParenLoc = Lex.getLoc();
+  if (!EatIfPresent(lltok::rparen))
+    return error(ParenLoc, "expected ')'");
+  if (!Depth)
+    return error(DepthLoc, "flatten_deep depth must be non-zero");
+  return false;
+}
+
 bool LLParser::parseOptionalUWTableKind(UWTableKind &Kind) {
   Lex.Lex();
   Kind = UWTableKind::Default;

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -2077,6 +2077,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::DisableSanitizerInstrumentation;
   case bitc::ATTR_KIND_ELEMENTTYPE:
     return Attribute::ElementType;
+  case bitc::ATTR_KIND_FLATTEN_DEEP:
+    return Attribute::FlattenDeep;
   case bitc::ATTR_KIND_FNRETTHUNK_EXTERN:
     return Attribute::FnRetThunkExtern;
   case bitc::ATTR_KIND_INLINE_HINT:
@@ -2391,6 +2393,8 @@ Error BitcodeReader::parseAttributeGroupBlock() {
             B.addDereferenceableAttr(Record[++i]);
           else if (Kind == Attribute::DereferenceableOrNull)
             B.addDereferenceableOrNullAttr(Record[++i]);
+          else if (Kind == Attribute::FlattenDeep)
+            B.addFlattenDeepAttr(Record[++i]);
           else if (Kind == Attribute::AllocSize)
             B.addAllocSizeAttrFromRawRepr(Record[++i]);
           else if (Kind == Attribute::VScaleRange)

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -956,6 +956,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_CAPTURES;
   case Attribute::DeadOnReturn:
     return bitc::ATTR_KIND_DEAD_ON_RETURN;
+  case Attribute::FlattenDeep:
+    return bitc::ATTR_KIND_FLATTEN_DEEP;
   case Attribute::EndAttrKinds:
     llvm_unreachable("Can not encode end-attribute kinds marker.");
   case Attribute::None:

llvm/lib/IR/Attributes.cpp

@@ -574,6 +574,9 @@ std::string Attribute::getAsString(bool InAttrGrp) const {
   if (hasAttribute(Attribute::DereferenceableOrNull))
     return AttrWithBytesToString("dereferenceable_or_null");
 
+  if (hasAttribute(Attribute::FlattenDeep))
+    return AttrWithBytesToString("flatten_deep");
+
   if (hasAttribute(Attribute::AllocSize)) {
     unsigned ElemSize;
     std::optional<unsigned> NumElems;
@@ -2206,6 +2209,13 @@ AttrBuilder &AttrBuilder::addDereferenceableOrNullAttr(uint64_t Bytes) {
   return addRawIntAttr(Attribute::DereferenceableOrNull, Bytes);
 }
 
+AttrBuilder &AttrBuilder::addFlattenDeepAttr(uint64_t Depth) {
+  if (Depth == 0)
+    return *this;
+
+  return addRawIntAttr(Attribute::FlattenDeep, Depth);
+}
+
 AttrBuilder &
 AttrBuilder::addAllocSizeAttr(unsigned ElemSize,
                               const std::optional<unsigned> &NumElems) {

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -936,6 +936,7 @@ Function *CodeExtractor::constructFunctionDeclaration(
         continue;
       // Those attributes should be safe to propagate to the extracted function.
       case Attribute::AlwaysInline:
+      case Attribute::FlattenDeep:
       case Attribute::Cold:
       case Attribute::DisableSanitizerInstrumentation:
       case Attribute::FnRetThunkExtern:

llvm/test/Bitcode/attributes-flatten-deep.ll

@@ -0,0 +1,32 @@
+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
+; RUN: verify-uselistorder %s
+
+; Test that flatten_deep attribute with integer values is properly handled
+; in both attribute groups (flatten_deep=N syntax) and inline (flatten_deep(N) syntax)
+
+; Test inline syntax
+; CHECK: define void @test_inline() #0
+define void @test_inline() flatten_deep(5) {
+  ret void
+}
+
+; Test attribute group alone
+; CHECK: define void @test_group_alone() #1
+define void @test_group_alone() #1 {
+  ret void
+}
+
+; Test attribute group with other attributes
+; CHECK: define void @test_group_combined() #2
+define void @test_group_combined() #2 {
+  ret void
+}
+
+; CHECK: attributes #0 = { flatten_deep=5 }
+attributes #0 = { flatten_deep=5 }
+
+; CHECK: attributes #1 = { flatten_deep=3 }
+attributes #1 = { flatten_deep=3 }
+
+; CHECK: attributes #2 = { noinline nounwind flatten_deep=7 }
+attributes #2 = { noinline nounwind flatten_deep=7 }