Add extra parameters for CUDA kernel function

Clang Frontend
1

Hi everyone,

I want to add some extra parameters for all CUDA kernel functions (any function that has the global attribute) inside a CUDA source file. For example, I have a function called matrixAddCUDA which was originally declared like this:

global void
matrixAddCUDA(float *C, float *A, float *B, int w, int h)
{

// code inside the function…
}

I want to use the Rewriter to add some extra parameters so that the function will be like this:

global void
matrixAddCUDA(float *C, float *A, float *B, int w, int h, int extra1, int extra2, int extra3)
{

// code inside the function
}

My current implementation is:

#include
#include

#include “clang/AST/AST.h”
#include “clang/AST/ASTConsumer.h”
#include “clang/ASTMatchers/ASTMatchFinder.h”
#include “clang/ASTMatchers/ASTMatchers.h”
#include “clang/Frontend/CompilerInstance.h”
#include “clang/Frontend/FrontendActions.h”
#include “clang/Rewrite/Core/Rewriter.h”
#include “clang/Tooling/CommonOptionsParser.h”
#include “clang/Tooling/Tooling.h”
#include “llvm/Support/raw_ostream.h”

using namespace clang::tooling;
using namespace llvm;
using namespace clang;
using namespace clang::ast_matchers;

// Apply a custom category to all command-line options so that they are the
// only ones displayed.
static llvm::cl::OptionCategory MyToolCategory(“my-tool options”);

// CommonOptionsParser declares HelpMessage with a description of the common
// command-line options related to the compilation database and input files.
// It’s nice to have this help message in all tools.
static cl::extrahelp CommonHelp(CommonOptionsParser::HelpMessage);

// A help message for this specific tool can be added afterwards.
static cl::extrahelp MoreHelp(“\nMore help text…”);

class MyASTConsumer: public ASTConsumer {
public:
Rewriter &TheRewriter;
MyASTConsumer (Rewriter &R) : TheRewriter(R){}
void HandleTranslationUnit(ASTContext &Context) override {
// Run the matchers when we have the whole TU parsed.
Finder.matchAST(Context);

for(DeclContext::decl_iterator D = Context.getTranslationUnitDecl()->decls_begin(), DEnd = Context.getTranslationUnitDecl()->decls_end(); D!=DEnd; ++D){
FunctionDecl* FD;
if((FD=dyn_cast (*D)) != 0){
if(FD->hasAttr()){
std::cout << “found " << FD->getNameAsString() << " fun decl\n”;
TypeSourceInfo *TI = FD->getTypeSourceInfo();
TypeLoc TL = TI->getTypeLoc();
FunctionTypeLoc FTL = TL.getAsAdjusted();
TheRewriter.insertText(FTL.getRParenLoc(), “int extra1, int extra2, int extra3”, true, true);
}
}

}

}

private:
MatchFinder Finder;
};

// For each source file provided to the tool, a new FrontendAction is created.
class MyFrontendAction : public ASTFrontendAction {
public:
MyFrontendAction() {}
void EndSourceFileAction() override {
TheRewriter.getEditBuffer(TheRewriter.getSourceMgr().getMainFileID())
.write(llvm::outs());
}

std::unique_ptr CreateASTConsumer(CompilerInstance &CI,
StringRef file) override {
TheRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
return llvm::make_unique(TheRewriter);
}

private:
Rewriter TheRewriter;
};

int main(int argc, const char **argv) {
CommonOptionsParser OptionsParser(argc, argv, MyToolCategory);
ClangTool Tool(OptionsParser.getCompilations(),
OptionsParser.getSourcePathList());

return Tool.run(newFrontendActionFactory().get());
}

It can successfully print out the CUDA kernel function, but cannot add extra parameters. I think there must be some problems with the Rewriter, but I don’t know where.

Thanks in advance!

Patrick