diff --git a/enzyme/Enzyme/Clang/EnzymeClang.cpp b/enzyme/Enzyme/Clang/EnzymeClang.cpp index a9e2de0c..3c54be7f 100644 --- a/enzyme/Enzyme/Clang/EnzymeClang.cpp +++ b/enzyme/Enzyme/Clang/EnzymeClang.cpp @@ -55,13 +55,13 @@ extern llvm::cl::opt ReactantBackend; std::vector GlobalOptimizationRules; -struct ByrefGlobalInfo { +struct TesseraArgTypeGlobalInfo { unsigned idx; QualType type; SourceLocation loc; }; -static std::vector ByrefGlobals; +static std::vector TesseraArgTypeGlobals; template class EnzymeAction final : public clang::PluginASTAction { @@ -108,12 +108,13 @@ static void emitOptimizationRules(Sema &S, std::vector &Rules) { } } -static void emitByrefGlobals(Sema &S, std::vector &Globals) { +static void +emitTesseraArgTypeGlobals(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 &Id = AST.Idents.get("__tessera_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); @@ -242,7 +243,7 @@ class EnzymePlugin final : public clang::ASTConsumer { void HandleTranslationUnit(ASTContext &Context) override { Sema &S = CI.getSema(); emitOptimizationRules(S, GlobalOptimizationRules); - emitByrefGlobals(S, ByrefGlobals); + emitTesseraArgTypeGlobals(S, TesseraArgTypeGlobals); } }; @@ -630,21 +631,27 @@ handleTesseraOpAttribute(Sema &S, Decl *D, const ParsedAttr &Attr, return ParsedAttrInfo::AttributeNotApplied; } - // Scan for byref positions + // Scan for val_in, val_out, and val_inout argument positions StringRef opStr = Literal->getString(); + bool hasArgList = opStr.contains('('); StringRef argList = opStr.slice(opStr.find('(') + 1, opStr.find(')')); - SmallVector byrefPositions; + SmallVector positionsToLift; + unsigned numListedArgs = 0; if (!argList.trim().empty()) { SmallVector argParts; argList.split(argParts, ','); + numListedArgs = argParts.size(); for (auto [idx, arg] : llvm::enumerate(argParts)) { arg = arg.trim(); - if (arg.contains(":byref") || arg.contains(": byref")) - byrefPositions.push_back(idx); + auto pos = arg.find(":"); + StringRef marker = arg.substr(pos + 1); + marker = marker.trim(); + if (marker == "val_inout" || marker == "val_in" || marker == "val_out") + positionsToLift.push_back(idx); } } - // Emit a global for each byref parameter + // Emit a global for each marked parameter auto FD = cast(D); DeclContext *declCtx = D->getDeclContext(); for (auto tmpCtx = declCtx; tmpCtx; tmpCtx = tmpCtx->getParent()) { @@ -655,27 +662,61 @@ handleTesseraOpAttribute(Sema &S, Decl *D, const ParsedAttr &Attr, auto params = FD->parameters(); auto loc = FD->getLocation(); + // A non-static C++ member function receives the object as an implicit + // leading `this` pointer, which is argument 0 of the emitted LLVM function + // but is absent from FD->parameters(). Positions in the op string name the + // call's arguments, so `this` occupies position 0 and the explicit + // parameters shift over by one: + // + // struct Mat { + // [[tessera::op("mfem.mult(this:val_in, x, y:val_out)")]] + // void Mult(const Vec &x, Vec &y) const; + // }; + const auto *MD = dyn_cast(FD); + bool hasImplicitThis = MD && MD->isInstance(); + + // The lowering requires one entry in the arg list per function argument + // (`this` included). + unsigned numExpectedArgs = params.size() + (hasImplicitThis ? 1 : 0); + if (hasArgList && numListedArgs != numExpectedArgs) { + unsigned ID = S.getDiagnostics().getCustomDiagID( + DiagnosticsEngine::Warning, + "'%0' argument list names %1 argument(s) but %2 takes %3%4; positions " + "in the argument list must match the function's arguments one for one"); + S.Diag(Attr.getLoc(), ID) + << attrName << numListedArgs << FD << numExpectedArgs + << (hasImplicitThis ? " (counting the implicit 'this')" : ""); + } + static unsigned globalCounter = 0; - SmallVector byrefGlobalIndices; + SmallVector liftedArgGlobalIndices; + + for (unsigned idx : positionsToLift) { + QualType pointeeTy; + if (hasImplicitThis && idx == 0) { + // The pointee of `this` is the (possibly const-qualified) class type. + pointeeTy = MD->getThisType()->getPointeeType(); + } else { + unsigned paramIdx = idx - (hasImplicitThis ? 1 : 0); + if (paramIdx >= params.size()) + continue; + pointeeTy = params[paramIdx]->getType(); + if (pointeeTy->isPointerType() || pointeeTy->isReferenceType()) + pointeeTy = (pointeeTy->getPointeeType()); + } - 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()); + pointeeTy = pointeeTy.getUnqualifiedType(); unsigned thisIdx = globalCounter++; - byrefGlobalIndices.push_back(thisIdx); - ByrefGlobals.push_back({thisIdx, pointeeTy, loc}); + liftedArgGlobalIndices.push_back(thisIdx); + TesseraArgTypeGlobals.push_back({thisIdx, pointeeTy, loc}); } - // Build annotation string: "tessera_op=eigen.inv(x:byref, y):3,4" + // Build annotation string: "tessera_op=eigen.inv(x:val_in, y):3,4" std::string annotation = (attrName + "=" + opStr).str(); // Parse remaining args representing sizes of function parameters - for (auto [i, idx] : llvm::enumerate(byrefGlobalIndices)) { + for (auto [i, idx] : llvm::enumerate(liftedArgGlobalIndices)) { annotation += (i == 0 ? ":globals=" : ",") + std::to_string(idx); }