From 50aefb8d954a6bee26b1550d165ea7e830e4dfb1 Mon Sep 17 00:00:00 2001 From: jessicacotturone21 Date: Sun, 26 Apr 2026 00:51:06 -0500 Subject: [PATCH 1/2] Added pure_tessera_op and extra arguments for parameter sizes --- enzyme/Enzyme/Clang/EnzymeClang.cpp | 104 ++++++++++++++++++++++++++-- 1 file changed, 100 insertions(+), 4 deletions(-) diff --git a/enzyme/Enzyme/Clang/EnzymeClang.cpp b/enzyme/Enzyme/Clang/EnzymeClang.cpp index 64a488a4..26c5aa9f 100644 --- a/enzyme/Enzyme/Clang/EnzymeClang.cpp +++ b/enzyme/Enzyme/Clang/EnzymeClang.cpp @@ -505,7 +505,7 @@ static ParsedAttrInfoRegistry::Add struct TesseraOpAttrInfo : public ParsedAttrInfo { TesseraOpAttrInfo() { - OptArgs = 1; + OptArgs = 15; // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and // [[plugin::example]] supported. static constexpr Spelling S[] = { @@ -534,13 +534,15 @@ struct TesseraOpAttrInfo : public ParsedAttrInfo { AttrHandling handleDeclAttribute(Sema &S, Decl *D, const ParsedAttr &Attr) const override { - if (Attr.getNumArgs() != 1) { + if (Attr.getNumArgs() < 1) { unsigned ID = S.getDiagnostics().getCustomDiagID( DiagnosticsEngine::Error, - "'tessera_op' attribute requires a single string argument"); + "'tessera_op' attribute requires at least a string argument"); S.Diag(Attr.getLoc(), ID); return AttributeNotApplied; } + + // Parse the first arg as the op string auto *Arg0 = Attr.getArgAsExpr(0); StringLiteral *Literal = dyn_cast(Arg0->IgnoreParenCasts()); if (!Literal) { @@ -551,8 +553,25 @@ struct TesseraOpAttrInfo : public ParsedAttrInfo { return AttributeNotApplied; } + // Build annotation string: "tessera_op=eigen.inv(x:byref, y):xsize,ysize" + std::string annotation = ("tessera_op=" + Literal->getString()).str(); + + // Parse remaining args representing sizes of function parameters + for (unsigned i = 1; i < Attr.getNumArgs(); i++) { + auto *SizeArg = Attr.getArgAsExpr(i); + Expr::EvalResult result; + if (!SizeArg->EvaluateAsInt(result, S.Context)) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, "argument %0 to 'tessera_op' attribute " + "must be an integer constant"); + S.Diag(Attr.getLoc(), ID) << i; + return AttributeNotApplied; + } + annotation += (i == 1 ? ":" : ",") + std::to_string(result.Val.getInt().getLimitedValue()); + } + D->addAttr(AnnotateAttr::Create( - S.Context, ("tessera_op=" + Literal->getString()).str(), + S.Context, annotation, nullptr, 0, Attr.getRange())); return AttributeApplied; } @@ -561,6 +580,83 @@ struct TesseraOpAttrInfo : public ParsedAttrInfo { static ParsedAttrInfoRegistry::Add T1("tessera_op", ""); +struct PureTesseraOpAttrInfo : public ParsedAttrInfo { + PureTesseraOpAttrInfo() { + OptArgs = 15; + // GNU-style __attribute__(("example")) and C++/C2x-style [[example]] and + // [[plugin::example]] supported. + static constexpr Spelling S[] = { + {ParsedAttr::AS_GNU, "pure_tessera_op"}, +#if LLVM_VERSION_MAJOR > 17 + {ParsedAttr::AS_C23, "pure_tessera_op"}, +#else + {ParsedAttr::AS_C2x, "pure_tessera_op"}, +#endif + {ParsedAttr::AS_CXX11, "pure_tessera_op"}, + {ParsedAttr::AS_CXX11, "tessera::pure_op"} + }; + Spellings = S; + } + + bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr, + const Decl *D) const override { + // This attribute appertains to functions only. + if (!isa(D)) { + S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str) + << Attr << "functions"; + return false; + } + return true; + } + + AttrHandling handleDeclAttribute(Sema &S, Decl *D, + const ParsedAttr &Attr) const override { + if (Attr.getNumArgs() < 1) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, + "'pure_tessera_op' attribute requires at least a string argument"); + S.Diag(Attr.getLoc(), ID); + return AttributeNotApplied; + } + + // Parse the first arg as the op string + auto *Arg0 = Attr.getArgAsExpr(0); + StringLiteral *Literal = dyn_cast(Arg0->IgnoreParenCasts()); + if (!Literal) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, "first argument to 'pure_tessera_op' " + "attribute must be a string literal"); + S.Diag(Attr.getLoc(), ID); + return AttributeNotApplied; + } + + // Build annotation string: "pure_tessera_op=eigen.inv(x:byref, y):xsize,ysize" + std::string annotation = ("pure_tessera_op=" + Literal->getString()).str(); + + // Parse remaining args representing sizes of function parameters + for (unsigned i = 1; i < Attr.getNumArgs(); i++) { + auto *SizeArg = Attr.getArgAsExpr(i); + Expr::EvalResult result; + if (!SizeArg->EvaluateAsInt(result, S.Context)) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, "argument %0 to 'pure_tessera_op' attribute " + "must be an integer constant"); + S.Diag(Attr.getLoc(), ID) << i; + return AttributeNotApplied; + } + annotation += (i == 1 ? ":" : ",") + std::to_string(result.Val.getInt().getLimitedValue()); + } + + D->addAttr(AnnotateAttr::Create( + S.Context, annotation, + nullptr, 0, Attr.getRange())); + return AttributeApplied; + } +}; + +static ParsedAttrInfoRegistry::Add + T2("pure_tessera_op", ""); + struct EnzymeShouldRecomputeAttrInfo : public ParsedAttrInfo { EnzymeShouldRecomputeAttrInfo() { OptArgs = 1; From 2eb8cd9b4386fd2462bc1fc0dfe80fbc7847a36a Mon Sep 17 00:00:00 2001 From: jessicacotturone21 Date: Mon, 29 Jun 2026 18:42:16 -0500 Subject: [PATCH 2/2] Added globals to contain types of byref arguments --- enzyme/Enzyme/Clang/EnzymeClang.cpp | 188 ++++++++++++++++------------ 1 file changed, 105 insertions(+), 83 deletions(-) diff --git a/enzyme/Enzyme/Clang/EnzymeClang.cpp b/enzyme/Enzyme/Clang/EnzymeClang.cpp index 26c5aa9f..ca7cc1db 100644 --- a/enzyme/Enzyme/Clang/EnzymeClang.cpp +++ b/enzyme/Enzyme/Clang/EnzymeClang.cpp @@ -53,6 +53,14 @@ extern llvm::cl::opt ReactantBackend; std::vector GlobalOptimizationRules; +struct ByrefGlobalInfo { + unsigned idx; + QualType type; + SourceLocation loc; +}; + +static std::vector ByrefGlobals; + template class EnzymeAction final : public clang::PluginASTAction { protected: @@ -82,7 +90,7 @@ static void emitOptimizationRules(Sema &S, std::vector &Rules) { DeclContext *declCtx = AST.getTranslationUnitDecl(); // create global variable for each optimization string - for (int i = 0; i < Rules.size(); i++) { + for (size_t i = 0, e = Rules.size(); i != e; ++i) { auto &Id = AST.Idents.get("__tessera_optimize_rule_" + std::to_string(i)); auto VD = VarDecl::Create(AST, declCtx, loc, loc, &Id, AST.IntTy, nullptr, SC_Static); @@ -97,6 +105,22 @@ static void emitOptimizationRules(Sema &S, std::vector &Rules) { } } +static void emitByrefGlobals(Sema &S, std::vector &Globals) { + auto &AST = S.getASTContext(); + DeclContext *declCtx = AST.getTranslationUnitDecl(); + for (auto &info : Globals) { + auto &Id = + AST.Idents.get("__tessera_byref_arg_type_" + std::to_string(info.idx)); + auto *VD = VarDecl::Create(AST, declCtx, info.loc, info.loc, &Id, info.type, + nullptr, SC_Static); + VD->setImplicit(true); + VD->setInit(new (AST) ImplicitValueInitExpr(info.type)); + VD->addAttr(clang::UsedAttr::CreateImplicit(AST)); + declCtx->addDecl(VD); + S.getASTConsumer().HandleTopLevelDecl(DeclGroupRef(VD)); + } +} + void MakeGlobalOfFn(FunctionDecl *FD, CompilerInstance &CI) { // if (FD->isLateTemplateParsed()) return; // TODO save any type info into string like attribute @@ -204,13 +228,13 @@ class EnzymePlugin final : public clang::ASTConsumer { CI.getPreprocessor().getHeaderSearchInfo().AddSearchPath(DL, /*isAngled=*/true); } + ~EnzymePlugin() {} void HandleTranslationUnit(ASTContext &Context) override { Sema &S = CI.getSema(); emitOptimizationRules(S, GlobalOptimizationRules); + emitByrefGlobals(S, ByrefGlobals); } - - ~EnzymePlugin() {} }; // register the PluginASTAction in the registry. @@ -503,6 +527,82 @@ struct EnzymeFunctionLikeAttrInfo : public ParsedAttrInfo { static ParsedAttrInfoRegistry::Add X4("enzyme_function_like", ""); +static ParsedAttrInfo::AttrHandling +handleTesseraOpAttribute(Sema &S, Decl *D, const ParsedAttr &Attr, + StringRef attrName) { + if (Attr.getNumArgs() < 1) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, + "'%0' attribute requires at least a string argument"); + S.Diag(Attr.getLoc(), ID) << attrName; + return ParsedAttrInfo::AttributeNotApplied; + } + + // Parse the first arg as the op string + auto *Arg0 = Attr.getArgAsExpr(0); + StringLiteral *Literal = dyn_cast(Arg0->IgnoreParenCasts()); + if (!Literal) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Error, "first argument to '%0' " + "attribute must be a string literal"); + S.Diag(Attr.getLoc(), ID) << attrName; + return ParsedAttrInfo::AttributeNotApplied; + } + + // Scan for byref positions + StringRef opStr = Literal->getString(); + StringRef argList = opStr.slice(opStr.find('(') + 1, opStr.find(')')); + SmallVector byrefPositions; + if (!argList.trim().empty()) { + SmallVector argParts; + argList.split(argParts, ','); + for (auto [idx, arg] : llvm::enumerate(argParts)) { + arg = arg.trim(); + if (arg.contains(":byref") || arg.contains(": byref")) + byrefPositions.push_back(idx); + } + } + + // Emit a global for each byref parameter + auto FD = cast(D); + DeclContext *declCtx = D->getDeclContext(); + for (auto tmpCtx = declCtx; tmpCtx; tmpCtx = tmpCtx->getParent()) { + if (tmpCtx->isRecord()) { + declCtx = tmpCtx->getParent(); + } + } + auto params = FD->parameters(); + auto loc = FD->getLocation(); + + static unsigned globalCounter = 0; + SmallVector byrefGlobalIndices; + + for (unsigned idx : byrefPositions) { + if (idx >= params.size()) + continue; + auto *param = params[idx]; + auto pointeeTy = param->getType(); + if (pointeeTy->isPointerType() || pointeeTy->isReferenceType()) + pointeeTy = (pointeeTy->getPointeeType()); + + unsigned thisIdx = globalCounter++; + byrefGlobalIndices.push_back(thisIdx); + ByrefGlobals.push_back({thisIdx, pointeeTy, loc}); + } + + // Build annotation string: "tessera_op=eigen.inv(x:byref, y):3,4" + std::string annotation = (attrName + "=" + opStr).str(); + + // Parse remaining args representing sizes of function parameters + for (auto [i, idx] : llvm::enumerate(byrefGlobalIndices)) { + annotation += (i == 0 ? ":globals=" : ",") + std::to_string(idx); + } + + D->addAttr( + AnnotateAttr::Create(S.Context, annotation, nullptr, 0, Attr.getRange())); + return ParsedAttrInfo::AttributeApplied; +} + struct TesseraOpAttrInfo : public ParsedAttrInfo { TesseraOpAttrInfo() { OptArgs = 15; @@ -534,46 +634,7 @@ struct TesseraOpAttrInfo : public ParsedAttrInfo { AttrHandling handleDeclAttribute(Sema &S, Decl *D, const ParsedAttr &Attr) const override { - if (Attr.getNumArgs() < 1) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, - "'tessera_op' attribute requires at least a string argument"); - S.Diag(Attr.getLoc(), ID); - return AttributeNotApplied; - } - - // Parse the first arg as the op string - auto *Arg0 = Attr.getArgAsExpr(0); - StringLiteral *Literal = dyn_cast(Arg0->IgnoreParenCasts()); - if (!Literal) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, "first argument to 'tessera_op' " - "attribute must be a string literal"); - S.Diag(Attr.getLoc(), ID); - return AttributeNotApplied; - } - - // Build annotation string: "tessera_op=eigen.inv(x:byref, y):xsize,ysize" - std::string annotation = ("tessera_op=" + Literal->getString()).str(); - - // Parse remaining args representing sizes of function parameters - for (unsigned i = 1; i < Attr.getNumArgs(); i++) { - auto *SizeArg = Attr.getArgAsExpr(i); - Expr::EvalResult result; - if (!SizeArg->EvaluateAsInt(result, S.Context)) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, "argument %0 to 'tessera_op' attribute " - "must be an integer constant"); - S.Diag(Attr.getLoc(), ID) << i; - return AttributeNotApplied; - } - annotation += (i == 1 ? ":" : ",") + std::to_string(result.Val.getInt().getLimitedValue()); - } - - D->addAttr(AnnotateAttr::Create( - S.Context, annotation, - nullptr, 0, Attr.getRange())); - return AttributeApplied; + return handleTesseraOpAttribute(S, D, Attr, "tessera_op"); } }; @@ -611,46 +672,7 @@ struct PureTesseraOpAttrInfo : public ParsedAttrInfo { AttrHandling handleDeclAttribute(Sema &S, Decl *D, const ParsedAttr &Attr) const override { - if (Attr.getNumArgs() < 1) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, - "'pure_tessera_op' attribute requires at least a string argument"); - S.Diag(Attr.getLoc(), ID); - return AttributeNotApplied; - } - - // Parse the first arg as the op string - auto *Arg0 = Attr.getArgAsExpr(0); - StringLiteral *Literal = dyn_cast(Arg0->IgnoreParenCasts()); - if (!Literal) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, "first argument to 'pure_tessera_op' " - "attribute must be a string literal"); - S.Diag(Attr.getLoc(), ID); - return AttributeNotApplied; - } - - // Build annotation string: "pure_tessera_op=eigen.inv(x:byref, y):xsize,ysize" - std::string annotation = ("pure_tessera_op=" + Literal->getString()).str(); - - // Parse remaining args representing sizes of function parameters - for (unsigned i = 1; i < Attr.getNumArgs(); i++) { - auto *SizeArg = Attr.getArgAsExpr(i); - Expr::EvalResult result; - if (!SizeArg->EvaluateAsInt(result, S.Context)) { - unsigned ID = S.getDiagnostics().getCustomDiagID( - DiagnosticsEngine::Error, "argument %0 to 'pure_tessera_op' attribute " - "must be an integer constant"); - S.Diag(Attr.getLoc(), ID) << i; - return AttributeNotApplied; - } - annotation += (i == 1 ? ":" : ",") + std::to_string(result.Val.getInt().getLimitedValue()); - } - - D->addAttr(AnnotateAttr::Create( - S.Context, annotation, - nullptr, 0, Attr.getRange())); - return AttributeApplied; + return handleTesseraOpAttribute(S, D, Attr, "pure_tessera_op"); } };