25#include "clang/AST/Attr.h"
26#include "clang/AST/DeclGroup.h"
27#include "clang/AST/RecursiveASTVisitor.h"
28#include "clang/Basic/FileManager.h"
29#include "clang/Basic/MacroBuilder.h"
30#include "clang/Frontend/CompilerInstance.h"
31#include "clang/Frontend/FrontendAction.h"
32#include "clang/Frontend/FrontendPluginRegistry.h"
33#include "clang/Lex/HeaderSearch.h"
34#include "clang/Lex/Preprocessor.h"
35#include "clang/Lex/PreprocessorOptions.h"
36#include "clang/Sema/Sema.h"
37#include "clang/Sema/SemaDiagnostic.h"
41#include "bundled_includes.h"
45template <
typename ConsumerType>
48 std::unique_ptr<clang::ASTConsumer>
50 llvm::StringRef InFile)
override {
51 return std::unique_ptr<clang::ASTConsumer>(
new ConsumerType(CI));
55 const std::vector<std::string> &args)
override {
60 return AddBeforeMainAction;
78#if LLVM_VERSION_MAJOR >= 18
87bool isDefinitionInThisTU(
const Decl *D) {
88 if (
auto FD = dyn_cast<FunctionDecl>(D))
89 return FD->isThisDeclarationADefinition();
90 if (
auto VD = dyn_cast<VarDecl>(D))
91 return VD->isThisDeclarationADefinition() != VarDecl::DeclarationOnly;
98bool registerEnzymeGlobal(Sema &S, Decl *D, StringRef Kind) {
99 auto &AST = S.getASTContext();
103 DeclContext *declCtx = D->getDeclContext();
104 for (
auto tmpCtx = declCtx; tmpCtx; tmpCtx = tmpCtx->getParent()) {
105 if (tmpCtx->isRecord()) {
106 declCtx = tmpCtx->getParent();
110 auto VD = cast<ValueDecl>(D);
111 auto loc = D->getLocation();
112 auto T = VD->getType();
113 auto FT = AST.getPointerType(T);
115 auto DR = DeclRefExpr::Create(AST, NestedNameSpecifierLoc(), loc, VD,
false,
116 loc, T, ExprValueKind::VK_LValue, VD,
118 auto rval = ExprValueKind::VK_PRValue;
119 Expr *expr =
nullptr;
120 if (isa<FunctionDecl>(D)) {
122 ImplicitCastExpr::Create(AST, FT, CastKind::CK_FunctionToPointerDecay,
123 DR,
nullptr, rval, FPOptionsOverride());
125 expr = UnaryOperator::Create(AST, DR, UnaryOperatorKind::UO_AddrOf, FT,
126 rval, clang::ExprObjectKind ::OK_Ordinary, loc,
127 false, FPOptionsOverride());
131 if (expr->isValueDependent())
134 auto &Id = AST.Idents.get(
135 (StringRef(
"__enzyme_") + Kind +
"_autoreg_" + VD->getNameAsString())
137 auto V = VarDecl::Create(AST, declCtx, loc, loc, &Id, FT,
nullptr, SC_None);
138 V->setStorageClass(SC_PrivateExtern);
139 V->addAttr(clang::UsedAttr::CreateImplicit(AST));
141 S.MarkVariableReferenced(loc, V);
142 S.getASTConsumer().HandleTopLevelDecl(DeclGroupRef(V));
148StringRef enzymeRegistrationKind(StringRef Annotation,
const Decl *D) {
149 if (isa<FunctionDecl>(D))
150 return llvm::StringSwitch<StringRef>(Annotation)
151 .Case(
"enzyme_inactivefn",
"inactivefn")
152 .Case(
"enzyme_inactivenoblockfn",
"inactivenoblockfn")
153 .Case(
"enzyme_nofree",
"nofree")
154 .Case(
"enzyme_sparse_accumulate",
"sparse_accumulate")
157 return llvm::StringSwitch<StringRef>(Annotation)
158 .Case(
"enzyme_inactive",
"inactive_global")
159 .Case(
"enzyme_nofree",
"nofree")
167void registerEnzymeDeclIfNotDefinedHere(Sema &S, Decl *D) {
168 if (!isa<FunctionDecl>(D) && !isa<VarDecl>(D))
170 if (isDefinitionInThisTU(D))
172 for (
auto A : D->specific_attrs<AnnotateAttr>()) {
173 StringRef Kind = enzymeRegistrationKind(A->getAnnotation(), D);
175 registerEnzymeGlobal(S, D, Kind);
182 clang::CompilerInstance &CI;
187 FrontendOptions &Opts = CI.getFrontendOpts();
188 CodeGenOptions &CGOpts = CI.getCodeGenOpts();
189 auto PluginName =
"ClangEnzyme-" + std::to_string(LLVM_VERSION_MAJOR);
191#if LLVM_VERSION_MAJOR < 18
192 std::string pluginPath;
194 for (
auto P : Opts.Plugins)
195 if (
endsWith(llvm::sys::path::stem(P), PluginName)) {
196#if LLVM_VERSION_MAJOR < 18
199 for (
auto passPlugin : CGOpts.PassPlugins) {
200 if (
endsWith(llvm::sys::path::stem(passPlugin), PluginName)) {
208#if LLVM_VERSION_MAJOR >= 18
211 CGOpts.PassPlugins.push_back(pluginPath);
214 CI.getPreprocessorOpts().Includes.push_back(
"/enzyme/enzyme/version");
216 std::string PredefineBuffer;
217 PredefineBuffer.reserve(4080);
218 llvm::raw_string_ostream Predefines(PredefineBuffer);
219 Predefines << CI.getPreprocessor().getPredefines();
220 MacroBuilder Builder(Predefines);
221 Builder.defineMacro(
"ENZYME_VERSION_MAJOR",
222 std::to_string(ENZYME_VERSION_MAJOR));
223 Builder.defineMacro(
"ENZYME_VERSION_MINOR",
224 std::to_string(ENZYME_VERSION_MINOR));
225 Builder.defineMacro(
"ENZYME_VERSION_PATCH",
226 std::to_string(ENZYME_VERSION_PATCH));
227 CI.getPreprocessor().setPredefines(Predefines.str());
229 auto baseFS = &CI.getFileManager().getVirtualFileSystem();
230 llvm::vfs::OverlayFileSystem *fuseFS(
231 new llvm::vfs::OverlayFileSystem(baseFS));
232 IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> fs(
233 new llvm::vfs::InMemoryFileSystem());
243 time_t timer = mktime(&y2k);
244 for (
const auto &pair : include_headers) {
245 fs->addFile(StringRef(pair[0]), timer,
246 llvm::MemoryBuffer::getMemBuffer(
247 StringRef(pair[1]), StringRef(pair[0]),
251 fuseFS->pushOverlay(fs);
252 fuseFS->pushOverlay(baseFS);
253 CI.getFileManager().setVirtualFileSystem(fuseFS);
255 auto DE = CI.getFileManager().getDirectoryRef(
"/enzymeroot");
257 auto DL = DirectoryLookup(*DE, SrcMgr::C_User,
259 CI.getPreprocessor().getHeaderSearchInfo().AddSearchPath(DL,
265 using namespace clang;
266 DeclGroupRef::iterator it;
270 for (it = dg.begin(); it != dg.end(); ++it) {
275 registerEnzymeDeclIfNotDefinedHere(CI.getSema(), *it);
276 if (
auto FD = dyn_cast<FunctionDecl>(*it)) {
277 if (!FD->hasAttr<clang::CUDADeviceAttr>())
280 if (!FD->getIdentifier())
282 if (!StringRef(FD->getLocation().printToString(CI.getSourceManager()))
283 .contains(
"/__clang_cuda_math.h") &&
284 !StringRef(FD->getLocation().printToString(CI.getSourceManager()))
285 .contains(
"/__clang_hip_math.h"))
288 FD->addAttr(UsedAttr::CreateImplicit(CI.getASTContext()));
290 if (
auto FD = dyn_cast<VarDecl>(*it)) {
297 if (!V->getIdentifier())
299 auto name = V->getName();
300 if (!(name.contains(
"__enzyme_inactive_global") ||
301 name.contains(
"__enzyme_inactivefn") ||
302 name.contains(
"__enzyme_inactivenoblockfn") ||
303 name.contains(
"__enzyme_shouldrecompute") ||
304 name.contains(
"__enzyme_function_like") ||
305 name.contains(
"__enzyme_allocation_like") ||
306 name.contains(
"__enzyme_register_gradient") ||
307 name.contains(
"__enzyme_register_derivative") ||
308 name.contains(
"__enzyme_register_splitderivative")))
311 V->addAttr(clang::UsedAttr::CreateImplicit(CI.getASTContext()));
317static clang::FrontendPluginRegistry::Add<EnzymeAction<EnzymePlugin>>
318 X(
"enzyme",
"Enzyme Plugin");
320#if LLVM_VERSION_MAJOR > 10
323bool isGlobalDecl(
const Decl *D) {
324 auto VD = dyn_cast<VarDecl>(D);
325 return VD && VD->hasGlobalStorage();
330bool appertainsToFunctionOrGlobal(Sema &S,
const ParsedAttr &Attr,
332 if (isa<FunctionDecl>(D) || isGlobalDecl(D))
335 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
336 << Attr <<
" applies to functions and globals only";
360ParsedAttrInfo::AttrHandling handleEnzymeMarkerAttr(Sema &S, Decl *D,
361 const ParsedAttr &Attr,
363 StringRef FnAnnotation,
364 StringRef VarAnnotation) {
365 if (Attr.getNumArgs() != 0) {
366 unsigned ID = S.getDiagnostics().getCustomDiagID(
367 DiagnosticsEngine::Error,
"'%0' attribute requires zero arguments");
368 S.Diag(Attr.getLoc(), ID) << AttrName;
369 return ParsedAttrInfo::AttributeNotApplied;
372 StringRef Annotation = isa<FunctionDecl>(D) ? FnAnnotation : VarAnnotation;
374 AnnotateAttr::Create(S.Context, Annotation,
nullptr, 0, Attr.getRange()));
376 if (D->getDeclContext()->isRecord())
377 registerEnzymeDeclIfNotDefinedHere(S, D);
378 return ParsedAttrInfo::AttributeApplied;
381struct EnzymeFunctionLikeAttrInfo :
public ParsedAttrInfo {
382 EnzymeFunctionLikeAttrInfo() {
386 static constexpr Spelling S[] = {
387 {ParsedAttr::AS_GNU,
"enzyme_function_like"},
388#if LLVM_VERSION_MAJOR > 17
389 {ParsedAttr::AS_C23,
"enzyme_function_like"},
391 {ParsedAttr::AS_C2x,
"enzyme_function_like"},
393 {ParsedAttr::AS_CXX11,
"enzyme_function_like"},
394 {ParsedAttr::AS_CXX11,
"enzyme::function_like"}
399 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
400 const Decl *D)
const override {
402 if (!isa<FunctionDecl>(D)) {
403 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
404 << Attr <<
"functions";
410 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
411 const ParsedAttr &Attr)
const override {
412 if (Attr.getNumArgs() != 1) {
413 unsigned ID = S.getDiagnostics().getCustomDiagID(
414 DiagnosticsEngine::Error,
415 "'enzyme_function' attribute requires a single string argument");
416 S.Diag(Attr.getLoc(), ID);
417 return AttributeNotApplied;
419 auto *Arg0 = Attr.getArgAsExpr(0);
420 StringLiteral *Literal = dyn_cast<StringLiteral>(Arg0->IgnoreParenCasts());
422 unsigned ID = S.getDiagnostics().getCustomDiagID(
423 DiagnosticsEngine::Error,
"first argument to 'enzyme_function_like' "
424 "attribute must be a string literal");
425 S.Diag(Attr.getLoc(), ID);
426 return AttributeNotApplied;
428 D->addAttr(AnnotateAttr::Create(
429 S.Context, (
"enzyme_function_like=" + Literal->getString()).str(),
430 nullptr, 0, Attr.getRange()));
431 return AttributeApplied;
435static ParsedAttrInfoRegistry::Add<EnzymeFunctionLikeAttrInfo>
436 X3(
"enzyme_function_like",
"");
438struct EnzymeShouldRecomputeAttrInfo :
public ParsedAttrInfo {
439 EnzymeShouldRecomputeAttrInfo() {
441 static constexpr Spelling S[] = {
442 {ParsedAttr::AS_GNU,
"enzyme_shouldrecompute"},
443#if LLVM_VERSION_MAJOR > 17
444 {ParsedAttr::AS_C23,
"enzyme_shouldrecompute"},
446 {ParsedAttr::AS_C2x,
"enzyme_shouldrecompute"},
448 {ParsedAttr::AS_CXX11,
"enzyme_shouldrecompute"},
449 {ParsedAttr::AS_CXX11,
"enzyme::shouldrecompute"}
454 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
455 const Decl *D)
const override {
457 if (isa<FunctionDecl>(D))
459 if (
auto VD = dyn_cast<VarDecl>(D)) {
460 if (VD->hasGlobalStorage())
463 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
464 << Attr <<
"functions and globals";
468 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
469 const ParsedAttr &Attr)
const override {
470 if (Attr.getNumArgs() != 0) {
471 unsigned ID = S.getDiagnostics().getCustomDiagID(
472 DiagnosticsEngine::Error,
473 "'enzyme_inactive' attribute requires zero arguments");
474 S.Diag(Attr.getLoc(), ID);
475 return AttributeNotApplied;
477 D->addAttr(AnnotateAttr::Create(S.Context,
"enzyme_shouldrecompute",
478 nullptr, 0, Attr.getRange()));
479 return AttributeApplied;
483static ParsedAttrInfoRegistry::Add<EnzymeShouldRecomputeAttrInfo>
484 ESR(
"enzyme_shouldrecompute",
"");
486struct EnzymeInactiveAttrInfo :
public ParsedAttrInfo {
487 EnzymeInactiveAttrInfo() {
491 static constexpr Spelling S[] = {
492 {ParsedAttr::AS_GNU,
"enzyme_inactive"},
493#if LLVM_VERSION_MAJOR > 17
494 {ParsedAttr::AS_C23,
"enzyme_inactive"},
496 {ParsedAttr::AS_C2x,
"enzyme_inactive"},
498 {ParsedAttr::AS_CXX11,
"enzyme_inactive"},
499 {ParsedAttr::AS_CXX11,
"enzyme::inactive"}
504 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
505 const Decl *D)
const override {
506 return appertainsToFunctionOrGlobal(S, Attr, D);
509 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
510 const ParsedAttr &Attr)
const override {
511 return handleEnzymeMarkerAttr(S, D, Attr,
"enzyme_inactive",
517static ParsedAttrInfoRegistry::Add<EnzymeInactiveAttrInfo> X4(
"enzyme_inactive",
520struct EnzymeInactiveNoblockAttrInfo :
public ParsedAttrInfo {
521 EnzymeInactiveNoblockAttrInfo() {
525 static constexpr Spelling S[] = {
526 {ParsedAttr::AS_GNU,
"enzyme_inactive_noblock"},
527#if LLVM_VERSION_MAJOR > 17
528 {ParsedAttr::AS_C23,
"enzyme_inactive_noblock"},
530 {ParsedAttr::AS_C2x,
"enzyme_inactive_noblock"},
532 {ParsedAttr::AS_CXX11,
"enzyme_inactive_noblock"},
533 {ParsedAttr::AS_CXX11,
"enzyme::inactive_noblock"}
538 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
539 const Decl *D)
const override {
540 return appertainsToFunctionOrGlobal(S, Attr, D);
543 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
544 const ParsedAttr &Attr)
const override {
545 return handleEnzymeMarkerAttr(S, D, Attr,
"enzyme_inactive_noblock",
546 "enzyme_inactivenoblockfn",
551static ParsedAttrInfoRegistry::Add<EnzymeInactiveNoblockAttrInfo>
552 X4_nb(
"enzyme_inactive_noblock",
"");
554struct EnzymeElementwiseReadAttrInfo :
public ParsedAttrInfo {
555 EnzymeElementwiseReadAttrInfo() {
557 static constexpr Spelling S[] = {
558 {ParsedAttr::AS_GNU,
"enzyme_elementwise_read"},
559#if LLVM_VERSION_MAJOR > 17
560 {ParsedAttr::AS_C23,
"enzyme_elementwise_read"},
562 {ParsedAttr::AS_C2x,
"enzyme_elementwise_read"},
564 {ParsedAttr::AS_CXX11,
"enzyme_elementwise_read"},
565 {ParsedAttr::AS_CXX11,
"enzyme::elementwise_read"}
570 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
571 const Decl *D)
const override {
572 if (isa<FunctionDecl>(D))
574 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
575 << Attr <<
"functions";
579 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
580 const ParsedAttr &Attr)
const override {
581 if (Attr.getNumArgs() != 0) {
582 unsigned ID = S.getDiagnostics().getCustomDiagID(
583 DiagnosticsEngine::Error,
584 "'enzyme_elementwise_read' attribute requires zero arguments");
585 S.Diag(Attr.getLoc(), ID);
586 return AttributeNotApplied;
588 D->addAttr(AnnotateAttr::Create(S.Context,
"enzyme_elementwise_read",
589 nullptr, 0, Attr.getRange()));
590 return AttributeApplied;
594static ParsedAttrInfoRegistry::Add<EnzymeElementwiseReadAttrInfo>
595 XElemRead(
"enzyme_elementwise_read",
"");
597struct EnzymeNoFreeAttrInfo :
public ParsedAttrInfo {
598 EnzymeNoFreeAttrInfo() {
602 static constexpr Spelling S[] = {
603 {ParsedAttr::AS_GNU,
"enzyme_nofree"},
604#if LLVM_VERSION_MAJOR > 17
605 {ParsedAttr::AS_C23,
"enzyme_nofree"},
607 {ParsedAttr::AS_C2x,
"enzyme_nofree"},
609 {ParsedAttr::AS_CXX11,
"enzyme_nofree"},
610 {ParsedAttr::AS_CXX11,
"enzyme::nofree"}
615 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
616 const Decl *D)
const override {
617 return appertainsToFunctionOrGlobal(S, Attr, D);
620 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
621 const ParsedAttr &Attr)
const override {
622 return handleEnzymeMarkerAttr(S, D, Attr,
"enzyme_nofree",
628static ParsedAttrInfoRegistry::Add<EnzymeNoFreeAttrInfo> X5(
"enzyme_nofree",
631struct EnzymeSparseAccumulateAttrInfo :
public ParsedAttrInfo {
632 EnzymeSparseAccumulateAttrInfo() {
636 static constexpr Spelling S[] = {
637 {ParsedAttr::AS_GNU,
"enzyme_sparse_accumulate"},
638#if LLVM_VERSION_MAJOR > 17
639 {ParsedAttr::AS_C23,
"enzyme_sparse_accumulate"},
641 {ParsedAttr::AS_C2x,
"enzyme_sparse_accumulate"},
643 {ParsedAttr::AS_CXX11,
"enzyme_sparse_accumulate"},
644 {ParsedAttr::AS_CXX11,
"enzyme::sparse_accumulate"}
649 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
650 const Decl *D)
const override {
652 if (isa<FunctionDecl>(D))
654 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
655 << Attr <<
"functions";
659 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
660 const ParsedAttr &Attr)
const override {
661 return handleEnzymeMarkerAttr(S, D, Attr,
"enzyme_sparse_accumulate",
662 "enzyme_sparse_accumulate",
663 "enzyme_sparse_accumulate");
667static ParsedAttrInfoRegistry::Add<EnzymeSparseAccumulateAttrInfo>
668 SparseX(
"enzyme_sparse_accumulate",
"");
670struct EnzymeNoTypeAnalysisAttrInfo :
public ParsedAttrInfo {
671 EnzymeNoTypeAnalysisAttrInfo() {
672 static constexpr Spelling S[] = {
673 {ParsedAttr::AS_GNU,
"enzyme_notypeanalysis"},
674#if LLVM_VERSION_MAJOR > 17
675 {ParsedAttr::AS_C23,
"enzyme_notypeanalysis"},
677 {ParsedAttr::AS_C2x,
"enzyme_notypeanalysis"},
679 {ParsedAttr::AS_CXX11,
"enzyme_notypeanalysis"},
680 {ParsedAttr::AS_CXX11,
"enzyme::notypeanalysis"}
685 bool diagAppertainsToDecl(Sema &S,
const ParsedAttr &Attr,
686 const Decl *D)
const override {
687 return appertainsToFunctionOrGlobal(S, Attr, D);
690 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
691 const ParsedAttr &Attr)
const override {
694 return handleEnzymeMarkerAttr(S, D, Attr,
"enzyme_notypeanalysis",
696 "enzyme_ta_norecur");
700static ParsedAttrInfoRegistry::Add<EnzymeNoTypeAnalysisAttrInfo>
701 enzyme_notypeanalysis(
"enzyme_notypeanalysis",
"");
static bool contains(ArrayRef< int > ar, int v)
static clang::FrontendPluginRegistry::Add< EnzymeAction< EnzymePlugin > > X("enzyme", "Enzyme Plugin")
void MakeGlobalOfFn(FunctionDecl *FD, CompilerInstance &CI)
void registerEnzyme(llvm::PassBuilder &PB)
static bool endsWith(llvm::StringRef string, llvm::StringRef suffix)
bool ParseArgs(const clang::CompilerInstance &CI, const std::vector< std::string > &args) override
std::unique_ptr< clang::ASTConsumer > CreateASTConsumer(clang::CompilerInstance &CI, llvm::StringRef InFile) override
PluginASTAction::ActionType getActionType() override
void HandleTranslationUnit(ASTContext &context) override
EnzymePlugin(clang::CompilerInstance &CI)
bool HandleTopLevelDecl(clang::DeclGroupRef dg) override
void HandleCXXStaticMemberVarInstantiation(clang::VarDecl *V) override
bool VisitFunctionDecl(FunctionDecl *FD)
Visitor(CompilerInstance &CI)