Enzyme main
Loading...
Searching...
No Matches
EnzymeClang.cpp
Go to the documentation of this file.
1//===- EnzymeClang.cpp - Automatic Differentiation Transformation Pass ----===//
2//
3// Enzyme Project
4//
5// Part of the Enzyme Project, under the Apache License v2.0 with LLVM
6// Exceptions. See https://llvm.org/LICENSE.txt for license information.
7// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
8//
9// If using this code in an academic setting, please cite the following:
10// @incollection{enzymeNeurips,
11// title = {Instead of Rewriting Foreign Code for Machine Learning,
12// Automatically Synthesize Fast Gradients},
13// author = {Moses, William S. and Churavy, Valentin},
14// booktitle = {Advances in Neural Information Processing Systems 33},
15// year = {2020},
16// note = {To appear in},
17// }
18//
19//===----------------------------------------------------------------------===//
20//
21// This file contains a clang plugin for Enzyme.
22//
23//===----------------------------------------------------------------------===//
24
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"
38
39#include "../Utils.h"
40
41#include "bundled_includes.h"
42
43using namespace clang;
44
45template <typename ConsumerType>
46class EnzymeAction final : public clang::PluginASTAction {
47protected:
48 std::unique_ptr<clang::ASTConsumer>
49 CreateASTConsumer(clang::CompilerInstance &CI,
50 llvm::StringRef InFile) override {
51 return std::unique_ptr<clang::ASTConsumer>(new ConsumerType(CI));
52 }
53
54 bool ParseArgs(const clang::CompilerInstance &CI,
55 const std::vector<std::string> &args) override {
56 return true;
57 }
58
59 PluginASTAction::ActionType getActionType() override {
60 return AddBeforeMainAction;
61 }
62};
63
64void MakeGlobalOfFn(FunctionDecl *FD, CompilerInstance &CI) {
65 // if (FD->isLateTemplateParsed()) return;
66 // TODO save any type info into string like attribute
67}
68
69struct Visitor : public RecursiveASTVisitor<Visitor> {
70 CompilerInstance &CI;
71 Visitor(CompilerInstance &CI) : CI(CI) {}
72 bool VisitFunctionDecl(FunctionDecl *FD) {
73 MakeGlobalOfFn(FD, CI);
74 return true;
75 }
76};
77
78#if LLVM_VERSION_MAJOR >= 18
79extern "C" void registerEnzyme(llvm::PassBuilder &PB);
80#endif
81
82namespace {
83
84/// Whether this declaration of D is a definition, and hence whether clang will
85/// emit an annotation on it -- annotations are only emitted for the entities a
86/// translation unit defines.
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;
92 return false;
93}
94
95/// Synthesize the global registering D with Enzyme, whose initializer is the
96/// address of D. Kind is the Enzyme registration kind, e.g. "inactivefn".
97/// Returns false if D is templated, and so has no address to register yet.
98bool registerEnzymeGlobal(Sema &S, Decl *D, StringRef Kind) {
99 auto &AST = S.getASTContext();
100
101 // The registration global is emitted alongside D, so step out of any
102 // enclosing record.
103 DeclContext *declCtx = D->getDeclContext();
104 for (auto tmpCtx = declCtx; tmpCtx; tmpCtx = tmpCtx->getParent()) {
105 if (tmpCtx->isRecord()) {
106 declCtx = tmpCtx->getParent();
107 }
108 }
109
110 auto VD = cast<ValueDecl>(D);
111 auto loc = D->getLocation();
112 auto T = VD->getType();
113 auto FT = AST.getPointerType(T);
114
115 auto DR = DeclRefExpr::Create(AST, NestedNameSpecifierLoc(), loc, VD, false,
116 loc, T, ExprValueKind::VK_LValue, VD,
117 /*TemplateArgs*/ nullptr);
118 auto rval = ExprValueKind::VK_PRValue;
119 Expr *expr = nullptr;
120 if (isa<FunctionDecl>(D)) {
121 expr =
122 ImplicitCastExpr::Create(AST, FT, CastKind::CK_FunctionToPointerDecay,
123 DR, nullptr, rval, FPOptionsOverride());
124 } else {
125 expr = UnaryOperator::Create(AST, DR, UnaryOperatorKind::UO_AddrOf, FT,
126 rval, clang::ExprObjectKind ::OK_Ordinary, loc,
127 /*canoverflow*/ false, FPOptionsOverride());
128 }
129
130 // A templated declaration has no address until it is instantiated.
131 if (expr->isValueDependent())
132 return false;
133
134 auto &Id = AST.Idents.get(
135 (StringRef("__enzyme_") + Kind + "_autoreg_" + VD->getNameAsString())
136 .str());
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));
140 V->setInit(expr);
141 S.MarkVariableReferenced(loc, V);
142 S.getASTConsumer().HandleTopLevelDecl(DeclGroupRef(V));
143 return true;
144}
145
146/// The Enzyme registration kind matching an annotation emitted by one of the
147/// marker attributes, or the empty string for any other annotation.
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")
155 .Default("");
156 if (isa<VarDecl>(D))
157 return llvm::StringSwitch<StringRef>(Annotation)
158 .Case("enzyme_inactive", "inactive_global")
159 .Case("enzyme_nofree", "nofree")
160 .Default("");
161 return "";
162}
163
164/// Clang only emits annotations for the entities a translation unit defines, so
165/// a declaration carrying one of the marker annotations but defined elsewhere
166/// needs a registration global to force the reference instead.
167void registerEnzymeDeclIfNotDefinedHere(Sema &S, Decl *D) {
168 if (!isa<FunctionDecl>(D) && !isa<VarDecl>(D))
169 return;
170 if (isDefinitionInThisTU(D))
171 return;
172 for (auto A : D->specific_attrs<AnnotateAttr>()) {
173 StringRef Kind = enzymeRegistrationKind(A->getAnnotation(), D);
174 if (!Kind.empty())
175 registerEnzymeGlobal(S, D, Kind);
176 }
177}
178
179} // namespace
180
181class EnzymePlugin final : public clang::ASTConsumer {
182 clang::CompilerInstance &CI;
183
184public:
185 EnzymePlugin(clang::CompilerInstance &CI) : CI(CI) {
186
187 FrontendOptions &Opts = CI.getFrontendOpts();
188 CodeGenOptions &CGOpts = CI.getCodeGenOpts();
189 auto PluginName = "ClangEnzyme-" + std::to_string(LLVM_VERSION_MAJOR);
190 bool contains = false;
191#if LLVM_VERSION_MAJOR < 18
192 std::string pluginPath;
193#endif
194 for (auto P : Opts.Plugins)
195 if (endsWith(llvm::sys::path::stem(P), PluginName)) {
196#if LLVM_VERSION_MAJOR < 18
197 pluginPath = P;
198#endif
199 for (auto passPlugin : CGOpts.PassPlugins) {
200 if (endsWith(llvm::sys::path::stem(passPlugin), PluginName)) {
201 contains = true;
202 break;
203 }
204 }
205 }
206
207 if (!contains) {
208#if LLVM_VERSION_MAJOR >= 18
209 CGOpts.PassBuilderCallbacks.push_back(registerEnzyme);
210#else
211 CGOpts.PassPlugins.push_back(pluginPath);
212#endif
213 }
214 CI.getPreprocessorOpts().Includes.push_back("/enzyme/enzyme/version");
215
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());
228
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());
234
235 struct tm y2k = {};
236
237 y2k.tm_hour = 0;
238 y2k.tm_min = 0;
239 y2k.tm_sec = 0;
240 y2k.tm_year = 100;
241 y2k.tm_mon = 0;
242 y2k.tm_mday = 1;
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]),
248 /*RequiresNullTerminator*/ true));
249 }
250
251 fuseFS->pushOverlay(fs);
252 fuseFS->pushOverlay(baseFS);
253 CI.getFileManager().setVirtualFileSystem(fuseFS);
254
255 auto DE = CI.getFileManager().getDirectoryRef("/enzymeroot");
256 assert(DE);
257 auto DL = DirectoryLookup(*DE, SrcMgr::C_User,
258 /*isFramework=*/false);
259 CI.getPreprocessor().getHeaderSearchInfo().AddSearchPath(DL,
260 /*isAngled=*/true);
261 }
263 void HandleTranslationUnit(ASTContext &context) override {}
264 bool HandleTopLevelDecl(clang::DeclGroupRef dg) override {
265 using namespace clang;
266 DeclGroupRef::iterator it;
267
268 // Visitor v(CI);
269 // Forcibly require emission of all libdevice
270 for (it = dg.begin(); it != dg.end(); ++it) {
271 // v.TraverseDecl(*it);
272 // Whether a declaration is defined here is only settled once it reaches
273 // the consumer, so this is where a marked declaration defined in another
274 // translation unit picks up its registration global.
275 registerEnzymeDeclIfNotDefinedHere(CI.getSema(), *it);
276 if (auto FD = dyn_cast<FunctionDecl>(*it)) {
277 if (!FD->hasAttr<clang::CUDADeviceAttr>())
278 continue;
279
280 if (!FD->getIdentifier())
281 continue;
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"))
286 continue;
287
288 FD->addAttr(UsedAttr::CreateImplicit(CI.getASTContext()));
289 }
290 if (auto FD = dyn_cast<VarDecl>(*it)) {
292 }
293 }
294 return true;
295 }
296 void HandleCXXStaticMemberVarInstantiation(clang::VarDecl *V) override {
297 if (!V->getIdentifier())
298 return;
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")))
309 return;
310
311 V->addAttr(clang::UsedAttr::CreateImplicit(CI.getASTContext()));
312 return;
313 }
314};
315
316// register the PluginASTAction in the registry.
317static clang::FrontendPluginRegistry::Add<EnzymeAction<EnzymePlugin>>
318 X("enzyme", "Enzyme Plugin");
319
320#if LLVM_VERSION_MAJOR > 10
321namespace {
322
323bool isGlobalDecl(const Decl *D) {
324 auto VD = dyn_cast<VarDecl>(D);
325 return VD && VD->hasGlobalStorage();
326}
327
328/// Shared check for the registration attributes which apply to both functions
329/// and global variables.
330bool appertainsToFunctionOrGlobal(Sema &S, const ParsedAttr &Attr,
331 const Decl *D) {
332 if (isa<FunctionDecl>(D) || isGlobalDecl(D))
333 return true;
334
335 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
336 << Attr << " applies to functions and globals only";
337 return false;
338}
339
340/// Shared implementation of the attributes which mark a declaration for Enzyme.
341///
342/// These lower to an annotation, which clang propagates through template
343/// instantiation for us, so the attribute also applies to declarations written
344/// inside a template -- where the alternative of taking the address of the
345/// declaration is not possible until instantiation. FnAnnotation is used for
346/// functions and VarAnnotation for variables.
347///
348/// Clang only emits annotations for the entities a translation unit defines,
349/// so a declaration defined elsewhere additionally needs a registration global
350/// to force the reference. Whether a function declarator will have a body is
351/// not yet known here, as attributes are processed before the body is parsed,
352/// so that is left to registerEnzymeDeclIfNotDefinedHere once the declaration
353/// reaches the AST consumer.
354///
355/// A class member never reaches the consumer on its own, so it is registered
356/// here instead, and for the same reason without knowing whether it is defined
357/// in this translation unit. A member defined here therefore gets both
358/// lowerings; they mark it identically, and the redundant global is consumed by
359/// PreserveNVVM like any other.
360ParsedAttrInfo::AttrHandling handleEnzymeMarkerAttr(Sema &S, Decl *D,
361 const ParsedAttr &Attr,
362 StringRef AttrName,
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;
370 }
371
372 StringRef Annotation = isa<FunctionDecl>(D) ? FnAnnotation : VarAnnotation;
373 D->addAttr(
374 AnnotateAttr::Create(S.Context, Annotation, nullptr, 0, Attr.getRange()));
375
376 if (D->getDeclContext()->isRecord())
377 registerEnzymeDeclIfNotDefinedHere(S, D);
378 return ParsedAttrInfo::AttributeApplied;
379}
380
381struct EnzymeFunctionLikeAttrInfo : public ParsedAttrInfo {
382 EnzymeFunctionLikeAttrInfo() {
383 OptArgs = 1;
384 // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and
385 // [[plugin::example]] supported.
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"},
390#else
391 {ParsedAttr::AS_C2x, "enzyme_function_like"},
392#endif
393 {ParsedAttr::AS_CXX11, "enzyme_function_like"},
394 {ParsedAttr::AS_CXX11, "enzyme::function_like"}
395 };
396 Spellings = S;
397 }
398
399 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
400 const Decl *D) const override {
401 // This attribute appertains to functions only.
402 if (!isa<FunctionDecl>(D)) {
403 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
404 << Attr << "functions";
405 return false;
406 }
407 return true;
408 }
409
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;
418 }
419 auto *Arg0 = Attr.getArgAsExpr(0);
420 StringLiteral *Literal = dyn_cast<StringLiteral>(Arg0->IgnoreParenCasts());
421 if (!Literal) {
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;
427 }
428 D->addAttr(AnnotateAttr::Create(
429 S.Context, ("enzyme_function_like=" + Literal->getString()).str(),
430 nullptr, 0, Attr.getRange()));
431 return AttributeApplied;
432 }
433};
434
435static ParsedAttrInfoRegistry::Add<EnzymeFunctionLikeAttrInfo>
436 X3("enzyme_function_like", "");
437
438struct EnzymeShouldRecomputeAttrInfo : public ParsedAttrInfo {
439 EnzymeShouldRecomputeAttrInfo() {
440 OptArgs = 1;
441 static constexpr Spelling S[] = {
442 {ParsedAttr::AS_GNU, "enzyme_shouldrecompute"},
443#if LLVM_VERSION_MAJOR > 17
444 {ParsedAttr::AS_C23, "enzyme_shouldrecompute"},
445#else
446 {ParsedAttr::AS_C2x, "enzyme_shouldrecompute"},
447#endif
448 {ParsedAttr::AS_CXX11, "enzyme_shouldrecompute"},
449 {ParsedAttr::AS_CXX11, "enzyme::shouldrecompute"}
450 };
451 Spellings = S;
452 }
453
454 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
455 const Decl *D) const override {
456 // This attribute appertains to functions only.
457 if (isa<FunctionDecl>(D))
458 return true;
459 if (auto VD = dyn_cast<VarDecl>(D)) {
460 if (VD->hasGlobalStorage())
461 return true;
462 }
463 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
464 << Attr << "functions and globals";
465 return false;
466 }
467
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;
476 }
477 D->addAttr(AnnotateAttr::Create(S.Context, "enzyme_shouldrecompute",
478 nullptr, 0, Attr.getRange()));
479 return AttributeApplied;
480 }
481};
482
483static ParsedAttrInfoRegistry::Add<EnzymeShouldRecomputeAttrInfo>
484 ESR("enzyme_shouldrecompute", "");
485
486struct EnzymeInactiveAttrInfo : public ParsedAttrInfo {
487 EnzymeInactiveAttrInfo() {
488 OptArgs = 1;
489 // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and
490 // [[plugin::example]] supported.
491 static constexpr Spelling S[] = {
492 {ParsedAttr::AS_GNU, "enzyme_inactive"},
493#if LLVM_VERSION_MAJOR > 17
494 {ParsedAttr::AS_C23, "enzyme_inactive"},
495#else
496 {ParsedAttr::AS_C2x, "enzyme_inactive"},
497#endif
498 {ParsedAttr::AS_CXX11, "enzyme_inactive"},
499 {ParsedAttr::AS_CXX11, "enzyme::inactive"}
500 };
501 Spellings = S;
502 }
503
504 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
505 const Decl *D) const override {
506 return appertainsToFunctionOrGlobal(S, Attr, D);
507 }
508
509 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
510 const ParsedAttr &Attr) const override {
511 return handleEnzymeMarkerAttr(S, D, Attr, "enzyme_inactive",
512 /*FnAnnotation*/ "enzyme_inactivefn",
513 /*VarAnnotation*/ "enzyme_inactive");
514 }
515};
516
517static ParsedAttrInfoRegistry::Add<EnzymeInactiveAttrInfo> X4("enzyme_inactive",
518 "");
519
520struct EnzymeInactiveNoblockAttrInfo : public ParsedAttrInfo {
521 EnzymeInactiveNoblockAttrInfo() {
522 OptArgs = 1;
523 // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and
524 // [[plugin::example]] supported.
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"},
529#else
530 {ParsedAttr::AS_C2x, "enzyme_inactive_noblock"},
531#endif
532 {ParsedAttr::AS_CXX11, "enzyme_inactive_noblock"},
533 {ParsedAttr::AS_CXX11, "enzyme::inactive_noblock"}
534 };
535 Spellings = S;
536 }
537
538 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
539 const Decl *D) const override {
540 return appertainsToFunctionOrGlobal(S, Attr, D);
541 }
542
543 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
544 const ParsedAttr &Attr) const override {
545 return handleEnzymeMarkerAttr(S, D, Attr, "enzyme_inactive_noblock",
546 /*FnAnnotation*/ "enzyme_inactivenoblockfn",
547 /*VarAnnotation*/ "enzyme_inactive");
548 }
549};
550
551static ParsedAttrInfoRegistry::Add<EnzymeInactiveNoblockAttrInfo>
552 X4_nb("enzyme_inactive_noblock", "");
553
554struct EnzymeElementwiseReadAttrInfo : public ParsedAttrInfo {
555 EnzymeElementwiseReadAttrInfo() {
556 OptArgs = 1;
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"},
561#else
562 {ParsedAttr::AS_C2x, "enzyme_elementwise_read"},
563#endif
564 {ParsedAttr::AS_CXX11, "enzyme_elementwise_read"},
565 {ParsedAttr::AS_CXX11, "enzyme::elementwise_read"}
566 };
567 Spellings = S;
568 }
569
570 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
571 const Decl *D) const override {
572 if (isa<FunctionDecl>(D))
573 return true;
574 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
575 << Attr << "functions";
576 return false;
577 }
578
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;
587 }
588 D->addAttr(AnnotateAttr::Create(S.Context, "enzyme_elementwise_read",
589 nullptr, 0, Attr.getRange()));
590 return AttributeApplied;
591 }
592};
593
594static ParsedAttrInfoRegistry::Add<EnzymeElementwiseReadAttrInfo>
595 XElemRead("enzyme_elementwise_read", "");
596
597struct EnzymeNoFreeAttrInfo : public ParsedAttrInfo {
598 EnzymeNoFreeAttrInfo() {
599 OptArgs = 1;
600 // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and
601 // [[plugin::example]] supported.
602 static constexpr Spelling S[] = {
603 {ParsedAttr::AS_GNU, "enzyme_nofree"},
604#if LLVM_VERSION_MAJOR > 17
605 {ParsedAttr::AS_C23, "enzyme_nofree"},
606#else
607 {ParsedAttr::AS_C2x, "enzyme_nofree"},
608#endif
609 {ParsedAttr::AS_CXX11, "enzyme_nofree"},
610 {ParsedAttr::AS_CXX11, "enzyme::nofree"}
611 };
612 Spellings = S;
613 }
614
615 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
616 const Decl *D) const override {
617 return appertainsToFunctionOrGlobal(S, Attr, D);
618 }
619
620 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
621 const ParsedAttr &Attr) const override {
622 return handleEnzymeMarkerAttr(S, D, Attr, "enzyme_nofree",
623 /*FnAnnotation*/ "enzyme_nofree",
624 /*VarAnnotation*/ "enzyme_nofree");
625 }
626};
627
628static ParsedAttrInfoRegistry::Add<EnzymeNoFreeAttrInfo> X5("enzyme_nofree",
629 "");
630
631struct EnzymeSparseAccumulateAttrInfo : public ParsedAttrInfo {
632 EnzymeSparseAccumulateAttrInfo() {
633 OptArgs = 1;
634 // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and
635 // [[plugin::example]] supported.
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"},
640#else
641 {ParsedAttr::AS_C2x, "enzyme_sparse_accumulate"},
642#endif
643 {ParsedAttr::AS_CXX11, "enzyme_sparse_accumulate"},
644 {ParsedAttr::AS_CXX11, "enzyme::sparse_accumulate"}
645 };
646 Spellings = S;
647 }
648
649 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
650 const Decl *D) const override {
651 // This attribute appertains to functions only.
652 if (isa<FunctionDecl>(D))
653 return true;
654 S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
655 << Attr << "functions";
656 return false;
657 }
658
659 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
660 const ParsedAttr &Attr) const override {
661 return handleEnzymeMarkerAttr(S, D, Attr, "enzyme_sparse_accumulate",
662 /*FnAnnotation*/ "enzyme_sparse_accumulate",
663 /*VarAnnotation*/ "enzyme_sparse_accumulate");
664 }
665};
666
667static ParsedAttrInfoRegistry::Add<EnzymeSparseAccumulateAttrInfo>
668 SparseX("enzyme_sparse_accumulate", "");
669
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"},
676#else
677 {ParsedAttr::AS_C2x, "enzyme_notypeanalysis"},
678#endif
679 {ParsedAttr::AS_CXX11, "enzyme_notypeanalysis"},
680 {ParsedAttr::AS_CXX11, "enzyme::notypeanalysis"}
681 };
682 Spellings = S;
683 }
684
685 bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
686 const Decl *D) const override {
687 return appertainsToFunctionOrGlobal(S, Attr, D);
688 }
689
690 AttrHandling handleDeclAttribute(Sema &S, Decl *D,
691 const ParsedAttr &Attr) const override {
692 // For now enzyme::notypeanalysis corresponds to the internal attribute
693 // enzyme_ta_norecur
694 return handleEnzymeMarkerAttr(S, D, Attr, "enzyme_notypeanalysis",
695 /*FnAnnotation*/ "enzyme_ta_norecur",
696 /*VarAnnotation*/ "enzyme_ta_norecur");
697 }
698};
699
700static ParsedAttrInfoRegistry::Add<EnzymeNoTypeAnalysisAttrInfo>
701 enzyme_notypeanalysis("enzyme_notypeanalysis", "");
702
703} // namespace
704
705#endif
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)
Definition Enzyme.cpp:3567
static bool endsWith(llvm::StringRef string, llvm::StringRef suffix)
Definition Utils.h:855
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)
CompilerInstance & CI