Add flatten_deep attribute in LLVM

Author: grigorypas

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2751,6 +2751,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/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{{.*}} }

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

@@ -802,6 +802,7 @@ enum AttributeKindCodes {
   ATTR_KIND_DEAD_ON_RETURN = 103,
   ATTR_KIND_SANITIZE_ALLOC_TOKEN = 104,
   ATTR_KIND_NO_CREATE_UNDEF_OR_POISON = 105,
+  ATTR_KIND_FLATTEN_DEEP = 106,
 };
 
 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

@@ -212,6 +212,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,23 @@ 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 +2511,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:
@@ -2393,6 +2395,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

@@ -958,6 +958,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_DEAD_ON_RETURN;
   case Attribute::NoCreateUndefOrPoison:
     return bitc::ATTR_KIND_NO_CREATE_UNDEF_OR_POISON;
+  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) {