Add extra parameters for CUDA kernel function

classic Classic list List threaded Threaded
1 message Options
Reply | Threaded
Open this post in threaded view
|

Add extra parameters for CUDA kernel function

Robinson, Paul via cfe-dev
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 <iostream>
#include <string>

#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<FunctionDecl> (*D)) != 0){
if(FD->hasAttr<CUDAGlobalAttr>()){ 
std::cout << "found " << FD->getNameAsString() << " fun decl\n";
TypeSourceInfo *TI = FD->getTypeSourceInfo();
TypeLoc TL = TI->getTypeLoc();
FunctionTypeLoc FTL = TL.getAsAdjusted<FunctionTypeLoc>();
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<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
StringRef file) override {
TheRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
return llvm::make_unique<MyASTConsumer>(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<MyFrontendAction>().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

_______________________________________________
cfe-dev mailing list
[hidden email]
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-dev