[clang] Support per-function [[clang::code_align(N)]] attribute.

The existing attribute works only for loop headers. This can
unfortunately be quite limiting, especially as a protection against the
"Jump Conditional Code Erratum" [0] (for example, if there is an
unaligned check in a hot loop). The command line option
-malign-branch-boundary can help with that, but it's too coarse-grained
and can significantly increase the binary size.

This change introduces a per-function [[clang::code_align(N)]] attribute
that aligns all the basic blocks of a function by the given N.

[0] https://www.intel.com/content/dam/support/us/en/documents/processors/mitigations-jump-conditional-code-erratum.pdf

Author: Anton Bikineev

clang/include/clang/Basic/Attr.td

@@ -4435,10 +4435,11 @@ def PreferredType: InheritableAttr {
   let Documentation = [PreferredTypeDocumentation];
 }
 
-def CodeAlign: StmtAttr {
+def CodeAlign : InheritableAttr {
   let Spellings = [Clang<"code_align">];
-  let Subjects = SubjectList<[ForStmt, CXXForRangeStmt, WhileStmt, DoStmt],
-                              ErrorDiag, "'for', 'while', and 'do' statements">;
+  let Subjects =
+      SubjectList<[ForStmt, CXXForRangeStmt, WhileStmt, DoStmt, Function],
+                  ErrorDiag, "'for', 'while', 'do' statements and functions">;
   let Args = [ExprArgument<"Alignment">];
   let Documentation = [CodeAlignAttrDocs];
   let AdditionalMembers = [{

clang/include/clang/Basic/AttrDocs.td

@@ -7704,11 +7704,12 @@ def CodeAlignAttrDocs : Documentation {
   let Category = DocCatVariable;
   let Heading = "clang::code_align";
   let Content = [{
-The ``clang::code_align(N)`` attribute applies to a loop and specifies the byte
-alignment for a loop. The attribute accepts a positive integer constant
-initialization expression indicating the number of bytes for the minimum
-alignment boundary. Its value must be a power of 2, between 1 and 4096
-(inclusive).
+The ``clang::code_align(N)`` attribute applies to a loop or a function. When
+applied to a loop it specifies the byte alignment for the loop header. When
+applied to a function it aligns all the basic blocks of the function. The
+attribute accepts a positive integer constant initialization expression
+indicating the number of bytes for the minimum alignment boundary. Its value
+must be a power of 2, between 1 and 4096 (inclusive).
 
 .. code-block:: c++
 
@@ -7738,6 +7739,11 @@ alignment boundary. Its value must be a power of 2, between 1 and 4096
     [[clang::code_align(A)]] for(;;) { }
   }
 
+  [[clang:code_align(16)]] int foo(bool b) {
+    if (b) return 2;
+    return 3;
+  }
+
   }];
 }
 

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2515,6 +2515,12 @@ void CodeGenModule::SetLLVMFunctionAttributesForDefinition(const Decl *D,
       F->setAlignment(std::max(llvm::Align(2), F->getAlign().valueOrOne()));
   }
 
+  if (const auto *CodeAlign = D->getAttr<CodeAlignAttr>()) {
+    const auto *CE = cast<ConstantExpr>(CodeAlign->getAlignment());
+    std::string ArgValStr = llvm::toString(CE->getResultAsAPSInt(), 10);
+    F->addFnAttr("align-basic-blocks", ArgValStr);
+  }
+
   // In the cross-dso CFI mode with canonical jump tables, we want !type
   // attributes on definitions only.
   if (CodeGenOpts.SanitizeCfiCrossDso &&

clang/lib/Sema/SemaDeclAttr.cpp

@@ -9036,6 +9036,12 @@ static void handleArmNewAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
                  ArmNewAttr(S.Context, AL, NewState.data(), NewState.size()));
 }
 
+static void handleCodeAlignAttr(Sema &S, Decl *D, const ParsedAttr &A) {
+  Expr *E = A.getArgAsExpr(0);
+  if (Attr *CodeAlignAttr = S.BuildCodeAlignAttr(A, E))
+    D->addAttr(CodeAlignAttr);
+}
+
 /// ProcessDeclAttribute - Apply the specific attribute to the specified decl if
 /// the attribute applies to decls.  If the attribute is a type attribute, just
 /// silently ignore it if a GNU attribute.
@@ -9855,6 +9861,10 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_UsingIfExists:
     handleSimpleAttribute<UsingIfExistsAttr>(S, D, AL);
     break;
+
+  case ParsedAttr::AT_CodeAlign:
+    handleCodeAlignAttr(S, D, AL);
+    break;
   }
 }
 

clang/test/CodeGen/code_align_function.c

@@ -0,0 +1,42 @@
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm -x c %s %s -o - | FileCheck -check-prefix=CHECK-C %s
+// RUN: %clang_cc1 -fsyntax-only -emit-llvm -x c++ -std=c++11 %s -o - | FileCheck %s --check-prefixes CHECK-CPP
+
+// CHECK-C: define dso_local i32 @code_align_function(i32 noundef %a) #[[A0:[0-9]+]]
+// CHECK-C: define dso_local i32 @code_align_function_with_aligned_loop(i32 noundef %a) #[[A1:[0-9]+]]
+
+// CHECK-C: attributes #[[A0]] = {{.*}} "align-basic-blocks"="32"
+[[clang::code_align(32)]] int code_align_function(int a) {
+  if (a) {
+    return 2;
+  }
+  return 3;
+}
+
+// CHECK-C: attributes #[[A1]] = {{.*}} "align-basic-blocks"="64"
+[[clang::code_align(64)]] int code_align_function_with_aligned_loop(int a) {
+  if (a) {
+    return 2;
+  }
+  int c = 0;
+  // CHECK-C: !{!"llvm.loop.align", i32 128}
+  [[clang::code_align(128)]] for (int i = 0; i < a; ++i) {
+    c += i;
+  }
+  return c;
+}
+
+#if __cplusplus >= 201103L
+struct S {
+  // CHECK-CPP: @_ZN1S19code_align_functionILi1EEEii({{.*}}) #[[A2:[0-9]+]]
+  // CHECK-CPP: attributes #[[A2]] = {{.*}} "align-basic-blocks"="16"
+  template <int A>
+  [[clang::code_align(16)]] int code_align_function(int a) {
+    if (a) {
+      return 2;
+    }
+    return 3;
+  }
+};
+
+template int S::code_align_function<1>(int);
+#endif

clang/test/Sema/code_align.c

@@ -9,14 +9,14 @@ void foo() {
   for (i = 0; i < 10; ++i) {  // this is OK
     a[i] = b[i] = 0;
   }
-  // expected-error@+1{{'code_align' attribute only applies to 'for', 'while', and 'do' statements}}
+  // expected-error@+1{{'code_align' attribute only applies to 'for', 'while', 'do' statements and functions}}
   [[clang::code_align(4)]]
   i = 7;
   for (i = 0; i < 10; ++i) {
     a[i] = b[i] = 0;
   }
 
-  // expected-error@+1{{'code_align' attribute cannot be applied to a declaration}}
+  // expected-error@+1{{'code_align' attribute only applies to 'for', 'while', 'do' statements and functions}}
   [[clang::code_align(12)]] int n[10];
 }
 
@@ -121,6 +121,11 @@ void check_code_align_expression() {
 #endif
 }
 
+[[clang::code_align(32)]] int function_attribute();
+
+// expected-error@+1{{'code_align' attribute requires an integer argument which is a constant power of two}}
+[[clang::code_align(31)]] int function_attribute_misaligned();
+
 #if __cplusplus >= 201103L
 template <int A, int B, int C, int D, int E>
 void code_align_dependent() {
@@ -161,6 +166,11 @@ void bar4() {
   for(int I=0; I<128; ++I) { bar(I); }
 }
 
+struct S {
+  template <int A>
+  [[clang::code_align(32)]] int foo();
+};
+
 int main() {
   code_align_dependent<8, 23, 32, -10, 64>(); // #neg-instantiation
   bar3<4>();  // #temp-instantiation

llvm/lib/CodeGen/MachineBlockPlacement.cpp

@@ -2940,8 +2940,12 @@ void MachineBlockPlacement::alignBlocks() {
       }
     }
 
-    // Use max of the TLIAlign and MDAlign
-    const Align LoopAlign = std::max(TLIAlign, Align(MDAlign));
+    unsigned FunctionMBBAlign =
+        F->getFunction().getFnAttributeAsParsedInteger("align-basic-blocks", 1);
+
+    // Use max of the TLIAlign, MDAlign or the function-level alignment.
+    const Align LoopAlign =
+        std::max(std::max(TLIAlign, Align(MDAlign)), Align(FunctionMBBAlign));
     if (LoopAlign == 1)
       continue; // Don't care about loop alignment.
 
@@ -3475,28 +3479,39 @@ bool MachineBlockPlacement::runOnMachineFunction(MachineFunction &MF) {
   bool HasMaxBytesOverride =
       MaxBytesForAlignmentOverride.getNumOccurrences() > 0;
 
+  unsigned long long MBBAlignment =
+      MF.getFunction().getFnAttributeAsParsedInteger("align-basic-blocks", 1);
+
+  // Respect BB alignment that could already be set by llvm.loop.align in
+  // alignBlocks() above.
   if (AlignAllBlock)
     // Align all of the blocks in the function to a specific alignment.
     for (MachineBasicBlock &MBB : MF) {
+      unsigned MaxAlignment = std::max(1ULL << AlignAllBlock, MBBAlignment);
       if (HasMaxBytesOverride)
-        MBB.setAlignment(Align(1ULL << AlignAllBlock),
+        MBB.setAlignment(std::max(Align(MaxAlignment), MBB.getAlignment()),
                          MaxBytesForAlignmentOverride);
       else
-        MBB.setAlignment(Align(1ULL << AlignAllBlock));
+        MBB.setAlignment(std::max(Align(MaxAlignment), MBB.getAlignment()));
     }
   else if (AlignAllNonFallThruBlocks) {
     // Align all of the blocks that have no fall-through predecessors to a
     // specific alignment.
     for (auto MBI = std::next(MF.begin()), MBE = MF.end(); MBI != MBE; ++MBI) {
       auto LayoutPred = std::prev(MBI);
+      unsigned MaxAlignment = std::max(1ULL << AlignAllNonFallThruBlocks, MBBAlignment);
       if (!LayoutPred->isSuccessor(&*MBI)) {
         if (HasMaxBytesOverride)
-          MBI->setAlignment(Align(1ULL << AlignAllNonFallThruBlocks),
+          MBI->setAlignment(std::max(Align(MaxAlignment), MBI->getAlignment()),
                             MaxBytesForAlignmentOverride);
         else
-          MBI->setAlignment(Align(1ULL << AlignAllNonFallThruBlocks));
+          MBI->setAlignment(std::max(Align(MaxAlignment), MBI->getAlignment()));
       }
     }
+  } else if (MBBAlignment != 1) {
+    for (MachineBasicBlock &MBB : MF) {
+      MBB.setAlignment(std::max(Align(MBBAlignment), MBB.getAlignment()));
+    }
   }
   if (ViewBlockLayoutWithBFI != GVDT_None &&
       (ViewBlockFreqFuncName.empty() ||

llvm/test/CodeGen/X86/code-align-basic-blocks.ll

@@ -0,0 +1,29 @@
+; RUN: llc < %s -mtriple=x86_64-pc-linux-gnu | FileCheck %s -check-prefixes=CHECK
+
+declare void @bar();
+
+; CHECK-LABEL: foo:
+; CHECK-LABEL: # %bb.0:
+; CHECK: .p2align        5, 0x90
+; CHECK-LABEL: # %bb.1:
+; CHECK: .p2align        10, 0x90
+; CHECK-LABEL: .LBB0_2:
+; CHECK: .p2align        5, 0x90
+; CHECK-LABEL: # %bb.3:
+; CHECK: .p2align        5, 0x90
+
+define i32 @foo(i1 %1) "align-basic-blocks"="32" {
+  br i1 %1, label %7, label %3
+3:
+  %4 = phi i32 [ %5, %3 ], [ 0, %2 ]
+  call void @bar()
+  %5 = add nuw nsw i32 %4, 1
+  %6 = icmp eq i32 %5, 90
+  br i1 %6, label %7, label %3, !llvm.loop !0
+7:
+  %8 = phi i32 [ 2, %2 ], [ 3, %3 ]
+  ret i32 %8
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.align", i32 1024}