[SYCLomatic] Fix the user define rule for macro (#2023)

Signed-off-by: Jiang, Zhiwei <[email protected]>
oneapi-src · Jun 3, 2024 · bbd870b · bbd870b
1 parent 92f5e18
commit bbd870b
Show file tree

Hide file tree

Showing 7 changed files with 121 additions and 93 deletions.
diff --git a/clang/lib/DPCT/ASTTraversal.cpp b/clang/lib/DPCT/ASTTraversal.cpp
@@ -169,7 +169,18 @@ void IncludesCallbacks::insertCudaArchRepl(
   return;
 }
 
-bool IncludesCallbacks::ReplaceCuMacro(const Token &MacroNameTok) {
+std::shared_ptr<clang::dpct::ReplaceToken>
+generateReplacement(SourceLocation SL, MacroMigrationRule Rule) {
+  requestFeature(Rule.HelperFeature);
+  for (auto ItHeader = Rule.Includes.begin(); ItHeader != Rule.Includes.end();
+       ItHeader++) {
+    DpctGlobalInfo::getInstance().insertHeader(SL, *ItHeader);
+  }
+  return std::make_shared<ReplaceToken>(SL, std::move(Rule.Out));
+}
+
+bool IncludesCallbacks::ReplaceCuMacro(const Token &MacroNameTok,
+                                       MacroInfo *MI) {
   bool IsInAnalysisScope = isInAnalysisScope(MacroNameTok.getLocation());
   if (!IsInAnalysisScope) {
     return false;
@@ -178,14 +189,11 @@ bool IncludesCallbacks::ReplaceCuMacro(const Token &MacroNameTok) {
     return false;
   }
   std::string MacroName = MacroNameTok.getIdentifierInfo()->getName().str();
-  auto Iter = MapNames::MacrosMap.find(MacroName);
-  if (Iter != MapNames::MacrosMap.end()) {
-    std::string ReplacedMacroName = Iter->second;
-    auto Repl = std::make_shared<ReplaceToken>(MacroNameTok.getLocation(),
-                                               std::move(ReplacedMacroName));
+  auto Iter = MapNames::MacroRuleMap.find(MacroName);
+  if (Iter != MapNames::MacroRuleMap.end()) {
+    auto Repl = generateReplacement(MacroNameTok.getLocation(), Iter->second);
     if (MacroName == "__CUDA_ARCH__") {
       if (DpctGlobalInfo::getInstance().getContext().getLangOpts().CUDA) {
-        requestFeature(HelperFeatureEnum::device_ext);
         insertCudaArchRepl(Repl->getReplacement(DpctGlobalInfo::getContext()));
         return true;
       }
@@ -197,16 +205,28 @@ bool IncludesCallbacks::ReplaceCuMacro(const Token &MacroNameTok) {
       } else {
         return false;
       }
-    }
-    if (MacroName == "CUDART_VERSION" || MacroName == "__CUDART_API_VERSION" ||
-        MacroName == "CUDA_VERSION") {
+    } else if (MacroName == "CUDART_VERSION" ||
+               MacroName == "__CUDART_API_VERSION" ||
+               MacroName == "CUDA_VERSION") {
       // These two macros are defined by CUDA header file
       auto LocInfo = DpctGlobalInfo::getLocInfo(MacroNameTok.getLocation());
       auto Ver = clang::getCudaVersionPair(DpctGlobalInfo::getSDKVersion());
       DpctGlobalInfo::getInstance()
           .insertFile(LocInfo.first)
           ->setRTVersionValue(
               std::to_string(Ver.first * 1000 + Ver.second * 10));
+    } else if (MacroName == "NCCL_VERSION_CODE" && MI) {
+      auto LocInfo = DpctGlobalInfo::getLocInfo(MacroNameTok.getLocation());
+      DpctGlobalInfo::getInstance()
+          .insertFile(LocInfo.first)
+          ->setCCLVerValue(Lexer::getSourceText(
+                               CharSourceRange::getCharRange(
+                                   MI->getReplacementToken(0).getLocation(),
+                                   Lexer::getLocForEndOfToken(
+                                       MI->getReplacementToken(0).getLocation(),
+                                       0, SM, LangOptions())),
+                               SM, LangOptions())
+                               .str());
     }
     if (DpctGlobalInfo::getContext().getLangOpts().CUDA) {
       // These two macros are defined by CUDA compiler
@@ -258,16 +278,10 @@ void IncludesCallbacks::MacroDefined(const Token &MacroNameTok,
     if (!II)
       continue;
 
-    if (MapNames::MacrosMap.find(II->getName().str()) !=
-        MapNames::MacrosMap.end()) {
-      std::string ReplacedMacroName =
-          MapNames::MacrosMap.at(II->getName().str());
+    auto ItRule = MapNames::MacroRuleMap.find(II->getName().str());
+    if (ItRule != MapNames::MacroRuleMap.end()) {
       TransformSet.emplace_back(
-          new ReplaceToken(Iter->getLocation(), std::move(ReplacedMacroName)));
-      if (II->getName().str() == "__CUDA_ARCH__" ||
-          II->getName().str() == "__NVCC__") {
-        requestFeature(HelperFeatureEnum::device_ext);
-      }
+          generateReplacement(Iter->getLocation(), ItRule->second));
     }
 
     if (II->hasMacroDefinition() && (II->getName().str() == "__host__" ||
@@ -455,9 +469,9 @@ void IncludesCallbacks::MacroExpands(const Token &MacroNameTok,
   if (!IsInAnalysisScope) {
     return;
   }
-  
-  if (ReplaceCuMacro(MacroNameTok)){
-    return ;
+
+  if (ReplaceCuMacro(MacroNameTok, MI)) {
+    return;
   }
 
   // For the un-specialized struct, there is no AST for the extern function
@@ -537,31 +551,6 @@ void IncludesCallbacks::MacroExpands(const Token &MacroNameTok,
                       .str())));
   }
 
-  auto ItRule = MapNames::MacroRuleMap.find(Name.str());
-  if (ItRule != MapNames::MacroRuleMap.end()) {
-    if (Name == "NCCL_VERSION_CODE") {
-      auto LocInfo = DpctGlobalInfo::getLocInfo(MacroNameTok.getLocation());
-      DpctGlobalInfo::getInstance()
-          .insertFile(LocInfo.first)
-          ->setCCLVerValue(Lexer::getSourceText(
-                               CharSourceRange::getCharRange(
-                                   MI->getReplacementToken(0).getLocation(),
-                                   Lexer::getLocForEndOfToken(
-                                       MI->getReplacementToken(0).getLocation(),
-                                       0, SM, LangOptions())),
-                               SM, LangOptions())
-                               .str());
-    }
-    std::string OutStr = ItRule->second.Out;
-    TransformSet.emplace_back(
-        new ReplaceToken(Range.getBegin(), std::move(OutStr)));
-    requestFeature(ItRule->second.HelperFeature);
-    for (auto ItHeader = ItRule->second.Includes.begin();
-         ItHeader != ItRule->second.Includes.end(); ItHeader++) {
-      DpctGlobalInfo::getInstance().insertHeader(Range.getBegin(), *ItHeader);
-    }
-  }
-
   if (TKind == tok::identifier && Name == "CUDART_CB") {
 #ifdef _WIN32
     TransformSet.emplace_back(
@@ -754,11 +743,10 @@ void IncludesCallbacks::ReplaceCuMacro(SourceRange ConditionRange, IfType IT,
   Token Tok;
   if (!Lexer::getRawToken(End, Tok, SM, LangOptions()))
     Size = Size + Tok.getLength();
-  std::string E(BP, Size);
-  for (auto &MacroMap : MapNames::MacrosMap) {
+  const std::string E(BP, Size);
+  for (auto &MacroRule : MapNames::MacroRuleMap) {
     size_t Pos = 0;
-    std::string MacroName = MacroMap.first;
-    std::string ReplacedMacroName = MacroMap.second;
+    std::string MacroName = MacroRule.first;
 
     std::size_t Found = E.find(MacroName, Pos);
     if (Found != std::string::npos && MacroName == "__CUDA_ARCH__") {
@@ -829,12 +817,11 @@ void IncludesCallbacks::ReplaceCuMacro(SourceRange ConditionRange, IfType IT,
       SourceLocation IB = Begin.getLocWithOffset(Found);
       SourceLocation IE = IB.getLocWithOffset(MacroName.length());
       CharSourceRange InsertRange(SourceRange(IB, IE), false);
-      auto Repl =
-          std::make_shared<ReplaceInclude>(InsertRange, ReplacedMacroName);
+
+      auto Repl = generateReplacement(IB, MacroRule.second);
       if (MacroName == "__CUDA_ARCH__" &&
           DpctGlobalInfo::getInstance().getContext().getLangOpts().CUDA) {
         insertCudaArchRepl(Repl->getReplacement(DpctGlobalInfo::getContext()));
-        requestFeature(HelperFeatureEnum::device_ext);
       } else if ((MacroName != "__CUDACC__" ||
                   DpctGlobalInfo::getMacroDefines().count(MacroName)) &&
                  MacroName != "__CUDA_ARCH__") {

diff --git a/clang/lib/DPCT/ASTTraversal.h b/clang/lib/DPCT/ASTTraversal.h
@@ -63,7 +63,7 @@ class IncludesCallbacks : public PPCallbacks {
   void Ifndef(SourceLocation Loc, const Token &MacroNameTok,
               const MacroDefinition &MD) override;
   // TODO: implement one of this for each source language.
-  bool ReplaceCuMacro(const Token &MacroNameTok);
+  bool ReplaceCuMacro(const Token &MacroNameTok, MacroInfo *MI = nullptr);
   void ReplaceCuMacro(SourceRange ConditionRange, IfType IT,
                       SourceLocation IfLoc, SourceLocation ElifLoc);
   void Defined(const Token &MacroNameTok, const MacroDefinition &MD,

diff --git a/clang/lib/DPCT/GenMakefile.cpp b/clang/lib/DPCT/GenMakefile.cpp
@@ -234,8 +234,8 @@ static void getCompileInfo(
           Len = Pos - strlen("-D");
         }
         std::string MacroName = Option.substr(strlen("-D"), Len);
-        auto Iter = MapNames::MacrosMap.find(MacroName);
-        if (Iter != MapNames::MacrosMap.end())
+        auto Iter = MapNames::MacroRuleMap.find(MacroName);
+        if (Iter != MapNames::MacroRuleMap.end())
           // Skip macros defined in helper function header files
           continue;
         else

diff --git a/clang/lib/DPCT/MapNames.cpp b/clang/lib/DPCT/MapNames.cpp
@@ -3887,25 +3887,6 @@ const MapNames::MapTy MapNames::Dim3MemberNamesMap{
     // ...
 };
 
-const MapNames::MapTy MapNames::MacrosMap{
-    {"__CUDA_ARCH__", "DPCT_COMPATIBILITY_TEMP"}, /**/
-    {"__NVCC__", "SYCL_LANGUAGE_VERSION"},
-    {"__CUDACC__", "SYCL_LANGUAGE_VERSION"},
-    {"__DRIVER_TYPES_H__", "__DPCT_HPP__"},
-    {"__CUDA_RUNTIME_H__", "__DPCT_HPP__"},
-    {"CUDART_VERSION", "DPCT_COMPAT_RT_VERSION"},
-    {"__CUDART_API_VERSION", "DPCT_COMPAT_RT_VERSION"},
-    {"CUDA_VERSION", "DPCT_COMPAT_RT_VERSION"},
-    {"__CUDACC_VER_MAJOR__", "DPCT_COMPAT_RT_MAJOR_VERSION"},
-    {"__CUDACC_VER_MINOR__", "DPCT_COMPAT_RT_MINOR_VERSION"},
-    {"CUBLAS_V2_H_", "MKL_SYCL_HPP"},
-    {"__CUDA__", "SYCL_LANGUAGE_VERSION"},
-    {"CUFFT_FORWARD", "-1"},
-    {"CUFFT_INVERSE", "1"},
-    {"cudaEventDefault", "0"},
-    //...
-};
-
 std::unordered_map<std::string, MacroMigrationRule> MapNames::MacroRuleMap{
     {"__forceinline__",
      MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
@@ -3924,39 +3905,90 @@ std::unordered_map<std::string, MacroMigrationRule> MapNames::MacroRuleMap{
                         "__noinline__", "__dpct_noinline__",
                         HelperFeatureEnum::device_ext)},
     {"cudaMemAttachGlobal",
-     MacroMigrationRule("flag_macro_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "cudaMemAttachGlobal", "0")},
     {"cudaStreamDefault",
-     MacroMigrationRule("cudaStreamDefault_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "cudaStreamDefault", "0")},
 
     {"CU_LAUNCH_PARAM_BUFFER_SIZE",
-     MacroMigrationRule("kernel_param_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "CU_LAUNCH_PARAM_BUFFER_SIZE", "((void *) 2)",
                         HelperFeatureEnum::device_ext)},
     {"CU_LAUNCH_PARAM_BUFFER_POINTER",
-     MacroMigrationRule("kernel_param_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "CU_LAUNCH_PARAM_BUFFER_POINTER", "((void *) 1)",
                         HelperFeatureEnum::device_ext)},
     {"CU_LAUNCH_PARAM_END",
-     MacroMigrationRule("kernel_param_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "CU_LAUNCH_PARAM_END", "((void *) 0)",
                         HelperFeatureEnum::device_ext)},
-    {"CUDART_PI_F", MacroMigrationRule("CUDART_PI_F", RulePriority::Fallback,
-                                       "CUDART_PI_F", "3.141592654F")},
-    {"CUB_MAX", MacroMigrationRule("cub_macro_rule", RulePriority::Fallback,
-                                   "CUB_MAX", "std::max")},
-    {"CUB_MIN", MacroMigrationRule("cub_macro_rule", RulePriority::Fallback,
-                                   "CUB_MIN", "std::min")},
+    {"CUDART_PI_F",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUDART_PI_F", "3.141592654F")},
+    {"CUB_MAX",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUB_MAX", "std::max")},
+    {"CUB_MIN",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUB_MIN", "std::min")},
     {"CUB_RUNTIME_FUNCTION",
-     MacroMigrationRule("cub_macro_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "CUB_RUNTIME_FUNCTION", "")},
     {"cudaStreamAttrValue",
-     MacroMigrationRule("cudaStreamAttrValue_macro_rule",
-                        RulePriority::Fallback, "cudaStreamAttrValue", "int")},
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "cudaStreamAttrValue", "int")},
     {"NCCL_VERSION_CODE",
-     MacroMigrationRule("nccl_macro_rule", RulePriority::Fallback,
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
                         "NCCL_VERSION_CODE", "DPCT_COMPAT_CCL_VERSION")},
+    {"__CUDA_ARCH__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDA_ARCH__", "DPCT_COMPATIBILITY_TEMP",
+                        clang::dpct::HelperFeatureEnum::device_ext)},
+    {"__NVCC__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__NVCC__", "SYCL_LANGUAGE_VERSION")},
+    {"__CUDACC__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDACC__", "SYCL_LANGUAGE_VERSION")},
+    {"__DRIVER_TYPES_H__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__DRIVER_TYPES_H__", "__DPCT_HPP__")},
+    {"__CUDA_RUNTIME_H__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDA_RUNTIME_H__", "__DPCT_HPP__")},
+    {"CUDART_VERSION",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUDART_VERSION", "DPCT_COMPAT_RT_VERSION")},
+    {"__CUDART_API_VERSION",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDART_API_VERSION", "DPCT_COMPAT_RT_VERSION")},
+    {"CUDA_VERSION",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUDA_VERSION", "DPCT_COMPAT_RT_VERSION")},
+    {"__CUDACC_VER_MAJOR__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDACC_VER_MAJOR__",
+                        "DPCT_COMPAT_RT_MAJOR_VERSION")},
+    {"__CUDACC_VER_MINOR__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDACC_VER_MINOR__",
+                        "DPCT_COMPAT_RT_MINOR_VERSION")},
+    {"CUBLAS_V2_H_",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUBLAS_V2_H_", "MKL_SYCL_HPP")},
+    {"__CUDA__",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "__CUDA__", "SYCL_LANGUAGE_VERSION")},
+    {"CUFFT_FORWARD",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUFFT_FORWARD", "-1")},
+    {"CUFFT_INVERSE",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "CUFFT_INVERSE", "1")},
+    {"cudaEventDefault",
+     MacroMigrationRule("dpct_build_in_macro_rule", RulePriority::Fallback,
+                        "cudaEventDefault", "0")},
     //...
 };
 

diff --git a/clang/lib/DPCT/MapNames.h b/clang/lib/DPCT/MapNames.h
@@ -324,7 +324,6 @@ class MapNames {
   static std::unordered_map<std::string, std::shared_ptr<TypeNameRule>>
       CuDNNTypeNamesMap;
   static const MapTy Dim3MemberNamesMap;
-  static const MapTy MacrosMap;
   static std::unordered_map<std::string, MacroMigrationRule> MacroRuleMap;
   static std::unordered_map<std::string, MetaRuleObject &> HeaderRuleMap;
   static MapTy BLASEnumsMap;

diff --git a/clang/test/dpct/user_defined_rule.cu b/clang/test/dpct/user_defined_rule.cu
@@ -199,5 +199,9 @@ void foo10(){
   filterfoo2(3);
 }
 
+// CHECK: #if defined(__UNDEFINED_MACRO__)
+// CHECK-NEXT: #endif
+#if defined(__NVCC__)
+#endif
 
 #endif
diff --git a/clang/test/dpct/user_defined_rule_2.yaml b/clang/test/dpct/user_defined_rule_2.yaml
@@ -4,4 +4,10 @@
   Priority: Takeover
   In: VECTOR
   Out: std::vector<int>
-  Includes: ["<vector>"]
+  Includes: ["<vector>"]
+- Rule: rule_macro_nvcc
+  Kind: Macro
+  Priority: Takeover
+  In: __NVCC__
+  Out: __UNDEFINED_MACRO__
+  Includes: []