[Transforms] Implement always_specialize attribute lowering

clang/include/clang/Basic/Attr.td

@@ -905,6 +905,14 @@ def AlwaysInline : DeclOrStmtAttr {
   let Documentation = [AlwaysInlineDocs];
 }
 
+def AlwaysSpecialize : InheritableParamAttr {
+ let Spellings = [GNU<"always_specialize">, CXX11<"clang", "always_specialize">,
+                   C23<"clang", "always_specialize">];
+ let Subjects = SubjectList<[ParmVar]>;
+ let Documentation = [AlwaysSpecializeDocs];
+ let SimpleHandler = 1;
+}
+
 def Artificial : InheritableAttr {
   let Spellings = [GCC<"artificial">];
   let Subjects = SubjectList<[InlineFunction]>;

clang/include/clang/Basic/AttrDocs.td

@@ -8126,6 +8126,17 @@ Attribute docs`_, and `the GCC Inline docs`_.
   let Heading = "always_inline, __force_inline";
 }
 
+def AlwaysSpecializeDocs : Documentation {
+  let Category = DocCatConsumed;
+  let Content = [{
+  The ``always_specialize`` attribute on a function parameter indicates that
+  the function shall be duplicated and specialized with respect to constant
+  arguments. This will usually increase code size. It controls an IR transform
+  similar in spirit to ``always_inline``.
+  }];
+    let Heading = "always_specialize";      
+}
+
 def EnforceTCBDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2559,6 +2559,9 @@ void CodeGenModule::SetLLVMFunctionAttributesForDefinition(const Decl *D,
       B.addAttribute("aarch64_new_zt0");
   }
 
+  if (D->hasAttr<AlwaysSpecializeAttr>())
+    B.addAttribute(llvm::Attribute::AlwaysSpecialize);
+
   // Track whether we need to add the optnone LLVM attribute,
   // starting with the default for this optimization level.
   bool ShouldAddOptNone =
@@ -2978,6 +2981,12 @@ void CodeGenModule::SetFunctionAttributes(GlobalDecl GD, llvm::Function *F,
     F->addParamAttr(0, llvm::Attribute::Returned);
   }
 
+  for (auto [Index, Param] : enumerate(FD->parameters()))
+    if (Param->hasAttrs())
+      for (auto *A : Param->getAttrs())
+        if (A->getKind() == attr::AlwaysSpecialize)
+          F->addParamAttr(Index, llvm::Attribute::AlwaysSpecialize);
+
   // Only a few attributes are set on declarations; these may later be
   // overridden by a definition.
 

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7137,6 +7137,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_AlwaysInline:
     handleAlwaysInlineAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_AlwaysSpecialize:
+    handleSimpleAttribute<AlwaysSpecializeAttr>(S, D, AL);
+    break;
   case ParsedAttr::AT_AnalyzerNoReturn:
     handleAnalyzerNoReturnAttr(S, D, AL);
     break;

clang/test/CodeGen/lto-newpm-pipeline.c

@@ -31,6 +31,7 @@
 // CHECK-FULL-O0-NEXT: Running pass: EntryExitInstrumenterPass
 // CHECK-FULL-O0-NEXT: Running pass: AlwaysInlinerPass
 // CHECK-FULL-O0-NEXT: Running analysis: ProfileSummaryAnalysis
+// CHECK-FULL-O0-NEXT: Running pass: AlwaysSpecializerPass
 // CHECK-FULL-O0-NEXT: Running pass: CoroConditionalWrapper
 // CHECK-FULL-O0-NEXT: Running pass: CanonicalizeAliasesPass
 // CHECK-FULL-O0-NEXT: Running pass: NameAnonGlobalPass
@@ -45,6 +46,7 @@
 // CHECK-THIN-O0-NEXT: Running pass: EntryExitInstrumenterPass
 // CHECK-THIN-O0-NEXT: Running pass: AlwaysInlinerPass
 // CHECK-THIN-O0-NEXT: Running analysis: ProfileSummaryAnalysis
+// CHECK-THIN-O0-NEXT: Running pass: AlwaysSpecializerPass
 // CHECK-THIN-O0-NEXT: Running pass: CoroConditionalWrapper
 // CHECK-THIN-O0-NEXT: Running pass: CanonicalizeAliasesPass
 // CHECK-THIN-O0-NEXT: Running pass: NameAnonGlobalPass

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

@@ -15,6 +15,7 @@
 // CHECK-NEXT: AlignValue (SubjectMatchRule_variable, SubjectMatchRule_type_alias)
 // CHECK-NEXT: AlwaysDestroy (SubjectMatchRule_variable)
 // CHECK-NEXT: AlwaysInline (SubjectMatchRule_function)
+// CHECK-NEXT: AlwaysSpecialize (SubjectMatchRule_variable_is_parameter)
 // CHECK-NEXT: Annotate ()
 // CHECK-NEXT: AnyX86NoCfCheck (SubjectMatchRule_hasType_functionType)
 // CHECK-NEXT: ArcWeakrefUnavailable (SubjectMatchRule_objc_interface)

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -798,6 +798,7 @@ enum AttributeKindCodes {
   ATTR_KIND_NO_DIVERGENCE_SOURCE = 100,
   ATTR_KIND_SANITIZE_TYPE = 101,
   ATTR_KIND_CAPTURES = 102,
+  ATTR_KIND_ALWAYS_SPECIALIZE = 103,
 };
 
 enum ComdatSelectionKindCodes {

llvm/include/llvm/IR/Attributes.td

@@ -98,6 +98,9 @@ def AllocSize : IntAttr<"allocsize", IntersectPreserve, [FnAttr]>;
 /// inline=always.
 def AlwaysInline : EnumAttr<"alwaysinline", IntersectPreserve, [FnAttr]>;
 
+/// Specialize function when argument at call site is known constant
+def AlwaysSpecialize : EnumAttr<"alwaysspecialize", IntersectPreserve, [ParamAttr]>;
+
 /// Callee is recognized as a builtin, despite nobuiltin attribute on its
 /// declaration.
 def Builtin : EnumAttr<"builtin", IntersectPreserve, [FnAttr]>;

llvm/include/llvm/InitializePasses.h

@@ -50,6 +50,7 @@ void initializeTarget(PassRegistry &);
 
 void initializeAAResultsWrapperPassPass(PassRegistry &);
 void initializeAlwaysInlinerLegacyPassPass(PassRegistry &);
+void initializeAlwaysSpecializerPass(PassRegistry &);
 void initializeAssignmentTrackingAnalysisPass(PassRegistry &);
 void initializeAssumptionCacheTrackerPass(PassRegistry &);
 void initializeAtomicExpandLegacyPass(PassRegistry &);

llvm/include/llvm/Transforms/IPO/AlwaysSpecializer.h

@@ -0,0 +1,29 @@
+//=== AlwaysSpecializer.h - implementation of always_specialize -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_IPO_ALWAYSSPECIALIZER_H
+#define LLVM_TRANSFORMS_IPO_ALWAYSSPECIALIZER_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class Module;
+class ModulePass;
+
+class AlwaysSpecializerPass : public PassInfoMixin<AlwaysSpecializerPass> {
+public:
+  AlwaysSpecializerPass();
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+  static bool isRequired() { return true; }
+};
+
+ModulePass *createAlwaysSpecializerPass();
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_ALWAYSSPECIALIZER_H

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -2054,6 +2054,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::Alignment;
   case bitc::ATTR_KIND_ALWAYS_INLINE:
     return Attribute::AlwaysInline;
+  case bitc::ATTR_KIND_ALWAYS_SPECIALIZE:
+    return Attribute::AlwaysSpecialize;
   case bitc::ATTR_KIND_BUILTIN:
     return Attribute::Builtin;
   case bitc::ATTR_KIND_BY_VAL:

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -750,6 +750,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_ALLOC_SIZE;
   case Attribute::AlwaysInline:
     return bitc::ATTR_KIND_ALWAYS_INLINE;
+  case Attribute::AlwaysSpecialize:
+    return bitc::ATTR_KIND_ALWAYS_SPECIALIZE;
   case Attribute::Builtin:
     return bitc::ATTR_KIND_BUILTIN;
   case Attribute::ByVal:

llvm/lib/Passes/PassBuilder.cpp

@@ -195,6 +195,7 @@
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "llvm/Transforms/HipStdPar/HipStdPar.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
+#include "llvm/Transforms/IPO/AlwaysSpecializer.h"
 #include "llvm/Transforms/IPO/Annotation2Metadata.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/IPO/Attributor.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -42,6 +42,7 @@
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "llvm/Transforms/HipStdPar/HipStdPar.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
+#include "llvm/Transforms/IPO/AlwaysSpecializer.h"
 #include "llvm/Transforms/IPO/Annotation2Metadata.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/IPO/Attributor.h"
@@ -1277,6 +1278,7 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
     MPM.addPass(PGOForceFunctionAttrsPass(PGOOpt->ColdOptType));
 
   MPM.addPass(AlwaysInlinerPass(/*InsertLifetimeIntrinsics=*/true));
+  MPM.addPass(AlwaysSpecializerPass());
 
   if (EnableModuleInliner)
     MPM.addPass(buildModuleInlinerPipeline(Level, Phase));
@@ -2252,6 +2254,7 @@ PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
   // code generation.
   MPM.addPass(AlwaysInlinerPass(
       /*InsertLifetimeIntrinsics=*/false));
+  MPM.addPass(AlwaysSpecializerPass());
 
   if (PTO.MergeFunctions)
     MPM.addPass(MergeFunctionsPass());

llvm/lib/Passes/PassRegistry.def

@@ -50,6 +50,7 @@ MODULE_ALIAS_ANALYSIS("globals-aa", GlobalsAA())
 #define MODULE_PASS(NAME, CREATE_PASS)
 #endif
 MODULE_PASS("always-inline", AlwaysInlinerPass())
+MODULE_PASS("always-specialize", AlwaysSpecializerPass())
 MODULE_PASS("annotation2metadata", Annotation2MetadataPass())
 MODULE_PASS("assign-guid", AssignGUIDPass())
 MODULE_PASS("attributor", AttributorPass())