diff --git a/hipify-clang/src/HipifyAction.cpp b/hipify-clang/src/HipifyAction.cpp index d6fe9b88e5..52ba2eeaaf 100644 --- a/hipify-clang/src/HipifyAction.cpp +++ b/hipify-clang/src/HipifyAction.cpp @@ -25,6 +25,7 @@ THE SOFTWARE. #include "clang/Frontend/CompilerInstance.h" #include "clang/ASTMatchers/ASTMatchFinder.h" #include "clang/ASTMatchers/ASTMatchers.h" +#include "clang/Lex/HeaderSearch.h" #include "LLVMCompat.h" #include "CUDA2HIP.h" #include "StringUtils.h" @@ -398,20 +399,50 @@ std::unique_ptr HipifyAction::CreateASTConsumer(clang::Compi return Finder->newASTConsumer(); } +void HipifyAction::Ifndef(clang::SourceLocation Loc, const clang::Token &MacroNameTok, const clang::MacroDefinition &MD) { + clang::SourceManager& SM = getCompilerInstance().getSourceManager(); + if (!SM.isWrittenInMainFile(Loc)) { + return; + } + StringRef Text(SM.getCharacterData(MacroNameTok.getLocation()), MacroNameTok.getLength()); + Ifndefs.insert(std::make_pair(Text.str(), MacroNameTok.getEndLoc())); +} + void HipifyAction::EndSourceFileAction() { // Insert the hip header, if we didn't already do it by accident during substitution. if (!insertedRuntimeHeader) { // It's not sufficient to just replace CUDA headers with hip ones, because numerous CUDA headers are // implicitly included by the compiler. Instead, we _delete_ CUDA headers, and unconditionally insert // one copy of the hip include into every file. + bool placeForIncludeCalculated = false; + clang::SourceLocation sl, controllingMacroLoc; clang::SourceManager& SM = getCompilerInstance().getSourceManager(); - clang::SourceLocation sl; + clang::Preprocessor& PP = getCompilerInstance().getPreprocessor(); + clang::HeaderSearch& HS = PP.getHeaderSearchInfo(); + clang::ExternalPreprocessorSource* EPL = HS.getExternalLookup(); + const clang::FileEntry* FE = SM.getFileEntryForID(SM.getMainFileID()); + const clang::IdentifierInfo* controllingMacro = HS.getFileInfo(FE).getControllingMacro(EPL); + if (controllingMacro) { + auto found = Ifndefs.find(controllingMacro->getName().str()); + if (found != Ifndefs.end()) { + controllingMacroLoc = found->second; + placeForIncludeCalculated = true; + } + } if (pragmaOnce) { - sl = pragmaOnceLoc; - } else if (firstHeader) { - sl = firstHeaderLoc; - } else { - sl = SM.getLocForStartOfFile(SM.getMainFileID()); + if (placeForIncludeCalculated) { + sl = pragmaOnceLoc < controllingMacroLoc ? pragmaOnceLoc : controllingMacroLoc; + } else { + sl = pragmaOnceLoc; + } + placeForIncludeCalculated = true; + } + if (!placeForIncludeCalculated) { + if (firstHeader) { + sl = firstHeaderLoc; + } else { + sl = SM.getLocForStartOfFile(SM.getMainFileID()); + } } clang::FullSourceLoc fullSL(sl, SM); ct::Replacement Rep(SM, sl, 0, "\n#include \n"); @@ -445,6 +476,10 @@ public: void PragmaDirective(clang::SourceLocation Loc, clang::PragmaIntroducerKind Introducer) override { hipifyAction.PragmaDirective(Loc, Introducer); } + + void Ifndef(clang::SourceLocation Loc, const clang::Token &MacroNameTok, const clang::MacroDefinition &MD) override { + hipifyAction.Ifndef(Loc, MacroNameTok, MD); + } }; } diff --git a/hipify-clang/src/HipifyAction.h b/hipify-clang/src/HipifyAction.h index cc7f3799ba..adafba2df9 100644 --- a/hipify-clang/src/HipifyAction.h +++ b/hipify-clang/src/HipifyAction.h @@ -40,6 +40,7 @@ class HipifyAction : public clang::ASTFrontendAction, public clang::ast_matchers::MatchFinder::MatchCallback { private: ct::Replacements* replacements; + std::map Ifndefs; std::unique_ptr Finder; // CUDA implicitly adds its runtime header. We rewrite explicitly-provided CUDA includes with equivalent // ones, and track - using this flag - if the result led to us including the hip runtime header. If it did @@ -81,6 +82,9 @@ public: const clang::Module *imported); // Called by the preprocessor for each pragma directive during the non-raw lexing pass. void PragmaDirective(clang::SourceLocation Loc, clang::PragmaIntroducerKind Introducer); + // Called by the preprocessor for each ifndef directive during the non-raw lexing pass. + // Found ifndef will be used in EndSourceFileAction() for catching include guard controlling macro. + void Ifndef(clang::SourceLocation Loc, const clang::Token &MacroNameTok, const clang::MacroDefinition &MD); protected: // Add a Replacement for the current file. These will all be applied after executing the FrontendAction. diff --git a/tests/hipify-clang/unit_tests/headers/headers_test_10.cu b/tests/hipify-clang/unit_tests/headers/headers_test_10.cu new file mode 100644 index 0000000000..1c2db50e3b --- /dev/null +++ b/tests/hipify-clang/unit_tests/headers/headers_test_10.cu @@ -0,0 +1,14 @@ +// RUN: %run_test hipify "%s" "%t" %hipify_args %clang_args + +// Checks that HIP header file is included after #pragma once, +// which goes before include guard controlling macro. +// CHECK: #pragma once +// CHECK-NEXT: #include +#pragma once +#ifndef HEADERS_TEST_10_H +// CHECK: #ifndef HEADERS_TEST_10_H +// CHECK-NOT: #include +#define HEADERS_TEST_10_H +#include +static int counter = 0; +#endif // HEADERS_TEST_10_H diff --git a/tests/hipify-clang/unit_tests/headers/headers_test_11.cu b/tests/hipify-clang/unit_tests/headers/headers_test_11.cu new file mode 100644 index 0000000000..7c59dbe4a7 --- /dev/null +++ b/tests/hipify-clang/unit_tests/headers/headers_test_11.cu @@ -0,0 +1,14 @@ +// RUN: %run_test hipify "%s" "%t" %hipify_args %clang_args + +// Checks that HIP header file is included after include guard controlling macro, +// which goes before #pragma once. +// CHECK: #ifndef HEADERS_TEST_10_H +// CHECK-NEXT: #include +#ifndef HEADERS_TEST_10_H +// CHECK: #pragma once +#pragma once +// CHECK-NOT: #include +#define HEADERS_TEST_10_H +#include +static int counter = 0; +#endif // HEADERS_TEST_10_H