[HIPIFY][fix][#211] Taking into account include guard controlling macro

...while including HIP main header file, which is inserted now after #indef controlling macro, or after #pragma once, if it's occurred earlier.

+ Add a couple of unit tests.
ToDo: Check backward compatibility on older clang versions.

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<clang::ASTConsumer> 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 <hip/hip_runtime.h>\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);
+  }
 };
 
 }

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<std::string, clang::SourceLocation> Ifndefs;
   std::unique_ptr<clang::ast_matchers::MatchFinder> 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.

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 <hip/hip_runtime.h>
+#pragma once
+#ifndef HEADERS_TEST_10_H
+// CHECK: #ifndef HEADERS_TEST_10_H
+// CHECK-NOT: #include <hip/hip_runtime.h>
+#define HEADERS_TEST_10_H
+#include <stdio.h>
+static int counter = 0;
+#endif // HEADERS_TEST_10_H

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 <hip/hip_runtime.h>
+#ifndef HEADERS_TEST_10_H
+// CHECK: #pragma once
+#pragma once
+// CHECK-NOT: #include <hip/hip_runtime.h>
+#define HEADERS_TEST_10_H
+#include <stdio.h>
+static int counter = 0;
+#endif // HEADERS_TEST_10_H