[SYCLomatic] Fix the device memory migration issue in the complex macro scenario (#2439)

Signed-off-by: intwanghao <[email protected]>

Author: intwanghao

clang/lib/DPCT/ASTTraversal.cpp

@@ -364,11 +364,25 @@ void IncludesCallbacks::MacroExpands(const Token &MacroNameTok,
       for (i = 0; i < MI->getNumTokens(); i++) {
         std::shared_ptr<dpct::DpctGlobalInfo::MacroExpansionRecord> R =
             std::make_shared<dpct::DpctGlobalInfo::MacroExpansionRecord>(
-                MacroNameTok.getIdentifierInfo(), MI, Range, IsInAnalysisScope, i);
+                MacroNameTok.getIdentifierInfo(), MI, Range, IsInAnalysisScope,
+                i);
         dpct::DpctGlobalInfo::getExpansionRangeToMacroRecord()
             [getCombinedStrFromLoc(MI->getReplacementToken(i).getLocation())] =
                 R;
       }
+      if (Args && IsInAnalysisScope) {
+        for (unsigned int i = 0; i < Args->getNumMacroArguments(); ++i) {
+          std::shared_ptr<dpct::DpctGlobalInfo::MacroArgRecord> R =
+              std::make_shared<dpct::DpctGlobalInfo::MacroArgRecord>(MI, i);
+          auto str =
+              getCombinedStrFromLoc(Args->getUnexpArgument(i)->getLocation());
+          auto &Global = DpctGlobalInfo::getInstance();
+          dpct::DpctGlobalInfo::getMacroArgRecordMap()
+              [Global.getMainFile()->getFilePath().getPath().str() +
+               getCombinedStrFromLoc(
+                   Args->getUnexpArgument(i)->getLocation())] = R;
+        }
+      }
       std::shared_ptr<dpct::DpctGlobalInfo::MacroExpansionRecord> R =
           std::make_shared<dpct::DpctGlobalInfo::MacroExpansionRecord>(
               MacroNameTok.getIdentifierInfo(), MI, Range, IsInAnalysisScope,
@@ -8634,13 +8648,6 @@ void MemVarRefMigrationRule::runRule(const MatchFinder::MatchResult &Result) {
     auto Info = Global.findMemVarInfo(Decl);
 
     if (Info && Info->isUseDeviceGlobal()) {
-      if (Decl->hasInit()) {
-        auto InitStr = getInitForDeviceGlobal(Decl);
-        if (!InitStr.empty()) {
-          report(Decl->getBeginLoc(), Diagnostics::DEVICE_GLOBAL_INIT, false);
-          Info->setInitForDeviceGlobal(InitStr);
-        }
-      }
       auto VarType = Info->getType();
       if (VarType->isArray()) {
         if (const auto *const ICE =
@@ -8786,13 +8793,8 @@ void ConstantMemVarMigrationRule::runRule(
     if (!Info)
       return;
     if (Info->isUseDeviceGlobal()) {
-      if (MemVar->hasInit()) {
-        auto InitStr = getInitForDeviceGlobal(MemVar);
-        if (!InitStr.empty()) {
-          report(MemVar->getBeginLoc(), Diagnostics::DEVICE_GLOBAL_INIT, false);
-          Info->setInitForDeviceGlobal(InitStr);
-        }
-      }
+      Info->migrateToDeviceGlobal(MemVar);
+      return;
     }
 
     Info->setIgnoreFlag(true);
@@ -9242,13 +9244,8 @@ void MemVarMigrationRule::runRule(
     if (!Info)
       return;
     if (Info->isUseDeviceGlobal()) {
-      if (MemVar->hasInit()) {
-        auto InitStr = getInitForDeviceGlobal(MemVar);
-        if (!InitStr.empty()) {
-          report(MemVar->getBeginLoc(), Diagnostics::DEVICE_GLOBAL_INIT, false);
-          Info->setInitForDeviceGlobal(InitStr);
-        }
-      }
+      Info->migrateToDeviceGlobal(MemVar);
+      return;
     }
 
     if (auto VTD = DpctGlobalInfo::findParent<VarTemplateDecl>(MemVar)) {

clang/lib/DPCT/AnalysisInfo.cpp

@@ -619,7 +619,8 @@ void DpctFileInfo::buildReplacements() {
   // found, postfix "_ct" is added to this __constant__ symbol's name.
   std::unordered_map<unsigned int, std::string> ReplUpdated;
   for (const auto &Entry : MemVarMap) {
-    if (Entry.second->isIgnore() || !Entry.second->isConstant())
+    if (Entry.second->isIgnore() || !Entry.second->isConstant() ||
+        Entry.second->isUseDeviceGlobal())
       continue;
 
     auto Name = Entry.second->getName();
@@ -1122,6 +1123,20 @@ DpctGlobalInfo::MacroDefRecord::MacroDefRecord(SourceLocation NTL, bool IIAS)
   FilePath = LocInfo.first;
   Offset = LocInfo.second;
 }
+
+DpctGlobalInfo::MacroArgRecord::MacroArgRecord(const MacroInfo *MI,
+                                               int ArgIndex)
+    : ArgIndex(ArgIndex) {
+  ArgName = MI->params()[ArgIndex]->getName().str();
+  for (auto Tok : MI->tokens()) {
+    auto II = Tok.getIdentifierInfo();
+    if (II && (II == MI->params()[ArgIndex])) {
+      ArgLoc = Tok.getLocation();
+      break;
+    }
+  }
+}
+
 DpctGlobalInfo::MacroExpansionRecord::MacroExpansionRecord(
     IdentifierInfo *ID, const MacroInfo *MI, SourceRange Range,
     bool IsInAnalysisScope, int TokenIndex) {
@@ -2399,6 +2414,8 @@ bool DpctGlobalInfo::CheckUnicodeSecurityFlag = false;
 bool DpctGlobalInfo::EnablepProfilingFlag = false;
 std::map<std::string, std::shared_ptr<DpctGlobalInfo::MacroExpansionRecord>>
     DpctGlobalInfo::ExpansionRangeToMacroRecord;
+std::unordered_map<std::string, std::shared_ptr<DpctGlobalInfo::MacroArgRecord>>
+    DpctGlobalInfo::MacroArgRecordMap;
 std::map<std::string, SourceLocation> DpctGlobalInfo::EndifLocationOfIfdef;
 std::vector<std::pair<clang::tooling::UnifiedPath, size_t>>
     DpctGlobalInfo::ConditionalCompilationLoc;
@@ -2831,6 +2848,163 @@ std::shared_ptr<MemVarInfo> MemVarInfo::buildMemVarInfo(const VarDecl *Var) {
   }
   return DpctGlobalInfo::getInstance().insertMemVarInfo(Var);
 }
+
+// This function, `migrateToDeviceGlobal`, migrates a CUDA `__device__` or
+// `__constant__` variable declaration to the SYCL device global equivalent. The
+// migration process involves four key steps. The function handles various
+// transformations as follows:
+//
+// 1. Remove any array brackets following the variable name.
+//    - It identifies the array type using TypeLoc and removes the brackets
+//    while preserving the dimensions for later use.
+//    - If the array size comes from a macro argument, it maps the macro
+//    argument correctly using the `MacroArgRecord`.
+// 2. Process the initialization expression.
+//    - If the initialization style is C-style (with an equal sign), it remove
+//    the equal sign and adds braces around scalar initializers to use
+//    initializer list in SYCL.
+// 3. Replace the variable type.
+//    - Replace the origin type with
+//    `sycl::ext::oneapi::experimental::device_global`
+//      and the correct base type and dimensions.
+//    - It manages macro arguments to correctly replace the base type when
+//    required.
+// 4. Insert the `static` specifier if the variable is declared globally and
+//    does not already have the `static` storage class.
+//
+// Example1 (Specifier __device__ will be removed in preprocessor callbacks):
+// Origin code:
+// __device__ int var_a[3] = {1, 2, 3};
+//
+// As follow list the result after each key step listed in previous:
+// 1. int var_a = {1, 2, 3};
+// 2. int var_a {1, 2, 3};
+// 3. sycl::ext::oneapi::experimental::device_global<int[3]> var_a {1, 2, 3};
+// 4. static sycl::ext::oneapi::experimental::device_global<int[3]> var_a {1, 2,
+// 3};
+//
+// Example2 (Specifier __device__ will be removed in preprocessor callbacks):
+// Origin code:
+// #define VAR(type, name, size) static __device__ type name[size];
+// VAR(int, a, 3)
+//
+// As follow list the result after each key step listed in previous:
+// 1. #define VAR(type, name, size) static type name;
+//    VAR(int, a, 3)
+// 2. #define VAR(type, name, init) static type name;
+//    VAR(int, a, 3)
+// 3. #define VAR(type, name, init) static
+//    sycl::ext::oneapi::experimental::device_global<type[size]> name;
+//    VAR(int, a, 3)
+// 4. #define VAR(type, name, init) static
+//    sycl::ext::oneapi::experimental::device_global<type[size]> name;
+//    VAR(int, a, 3)
+void MemVarInfo::migrateToDeviceGlobal(const VarDecl *MemVar) {
+  auto &SM = DpctGlobalInfo::getSourceManager();
+  auto &Ctx = DpctGlobalInfo::getContext();
+  auto &MacroArgMap = DpctGlobalInfo::getMacroArgRecordMap();
+  auto TSI = MemVar->getTypeSourceInfo();
+  auto OriginTL = TSI->getTypeLoc();
+  auto TL = OriginTL;
+  auto BegLoc = MemVar->getBeginLoc();
+  if (BegLoc.isMacroID()) {
+    BegLoc = SM.getExpansionLoc(BegLoc);
+  }
+  auto LocInfo = DpctGlobalInfo::getLocInfo(BegLoc);
+  std::string Dims;
+  bool IsArray = OriginTL.getType()->isArrayType();
+  // Step 1
+  while (auto ATL = TL.getAs<clang::ArrayTypeLoc>()) {
+    auto BRange = ATL.getBracketsRange();
+    BRange = getDefinitionRange(BRange.getBegin(), BRange.getEnd());
+    auto RT =
+        ReplaceText(SM.getSpellingLoc(BRange.getBegin()),
+                    SM.getSpellingLoc(BRange.getEnd()).getLocWithOffset(1), "");
+    DpctGlobalInfo::getInstance().addReplacement(RT.getReplacement(Ctx));
+    Dims += "[";
+    std::string SizeStr;
+    if (clang::Expr *SE = ATL.getSizeExpr()) {
+      auto SizeLoc = SE->getBeginLoc();
+      if (SM.isMacroArgExpansion(SizeLoc)) {
+        auto Iter =
+            MacroArgMap.find(DpctGlobalInfo::getInstance()
+                                 .getMainFile()
+                                 ->getFilePath()
+                                 .getPath()
+                                 .str() +
+                             getCombinedStrFromLoc(SM.getSpellingLoc(SizeLoc)));
+        if (Iter != MacroArgMap.end()) {
+          SizeStr = Iter->second->ArgName;
+        }
+      }
+      if (SizeStr.empty()) {
+        SizeStr = ExprAnalysis::ref(SE);
+      }
+    }
+    Dims += SizeStr + "]";
+    TL = ATL.getElementLoc();
+  }
+  // Step 2
+  if (MemVar->hasInit()) {
+    if ((MemVar->getInitStyle() == VarDecl::InitializationStyle::CInit)) {
+      DiagnosticsUtils::report(LocInfo.first, LocInfo.second,
+                               Diagnostics::DEVICE_GLOBAL_INIT, true, false);
+      if (!dyn_cast<InitListExpr>(
+              MemVar->getInit()->IgnoreImplicitAsWritten())) {
+        auto IBS = InsertBeforeStmt(MemVar->getInit(), "{");
+        auto IAS = InsertAfterStmt(MemVar->getInit(), "}");
+        DpctGlobalInfo::getInstance().addReplacement(IBS.getReplacement(Ctx));
+        DpctGlobalInfo::getInstance().addReplacement(IAS.getReplacement(Ctx));
+      }
+      auto NextTok = Lexer::findNextToken(
+          IsArray ? SM.getSpellingLoc(OriginTL.getEndLoc())
+                  : SM.getSpellingLoc(MemVar->getLocation()),
+          SM, DpctGlobalInfo::getContext().getLangOpts());
+      if (NextTok.has_value() && NextTok.value().is(tok::equal)) {
+        auto RTok = ReplaceToken(NextTok.value().getLocation(), "");
+        DpctGlobalInfo::getInstance().addReplacement(RTok.getReplacement(Ctx));
+      }
+    }
+  }
+  // Step 3
+  std::string BaseTypeStr;
+  SourceLocation TypeReplLoc;
+  size_t TypeReplLen = 0;
+  if (SM.isMacroArgExpansion(OriginTL.getBeginLoc())) {
+    auto Iter = MacroArgMap.find(
+        DpctGlobalInfo::getInstance()
+            .getMainFile()
+            ->getFilePath()
+            .getPath()
+            .str() +
+        getCombinedStrFromLoc(SM.getSpellingLoc(OriginTL.getBeginLoc())));
+    if (Iter != MacroArgMap.end()) {
+      BaseTypeStr = Iter->second->ArgName;
+      TypeReplLoc = Iter->second->ArgLoc;
+      TypeReplLen = BaseTypeStr.size();
+    }
+  }
+  if (BaseTypeStr.empty()) {
+    BaseTypeStr = getType()->getBaseNameWithoutQualifiers();
+    TypeReplLoc = TL.getBeginLoc();
+    TypeReplLen =
+        SM.getFileOffset(TL.getEndLoc()) - SM.getFileOffset(TL.getBeginLoc()) +
+        Lexer::MeasureTokenLength(TL.getEndLoc(), SM,
+                                  DpctGlobalInfo::getContext().getLangOpts());
+  }
+  std::string TypeStr = MapNames::getClNamespace() +
+                        "ext::oneapi::experimental::device_global<" +
+                        BaseTypeStr + Dims + ">";
+  auto RT = ReplaceText(TypeReplLoc, TypeReplLen, std::move(TypeStr));
+  DpctGlobalInfo::getInstance().addReplacement(RT.getReplacement(Ctx));
+  // Step 4
+  if (MemVar->getStorageClass() != SC_Static && getScope() == Global) {
+    DpctGlobalInfo::getInstance().addReplacement(
+        std::make_shared<ExtReplacement>(LocInfo.first, LocInfo.second, 0,
+                                         "static ", nullptr));
+  }
+}
+
 MemVarInfo::VarAttrKind MemVarInfo::getAddressAttr(const VarDecl *VD) {
   if (VD->hasAttrs())
     return getAddressAttr(VD->getAttrs());

clang/lib/DPCT/AnalysisInfo.h

@@ -633,6 +633,22 @@ class DpctGlobalInfo {
     bool IsInAnalysisScope;
     MacroDefRecord(SourceLocation NTL, bool IIAS);
   };
+  // This class is used to store information about macro arguments in a
+  // macro definition. For example, consider the macro definition:
+  // "#define CALL(x, y) x(y)".
+  // - For the first argument "x", the member ArgName will be "x", ArgLoc will
+  // be the source location of the token "x" in the macro definition, and
+  // ArgIndex will be 0.
+  // - For the second argument "y", the member ArgName will be "y", ArgLoc will
+  // be the source location of the token "y" in the macro definition, and
+  // ArgIndex will be 1.
+  class MacroArgRecord {
+  public:
+    std::string ArgName;
+    SourceLocation ArgLoc;
+    int ArgIndex;
+    MacroArgRecord(const MacroInfo *MI, int ArgIndex);
+  };
 
   class MacroExpansionRecord {
   public:
@@ -1189,6 +1205,10 @@ class DpctGlobalInfo {
   getExpansionRangeBeginMap() {
     return ExpansionRangeBeginMap;
   }
+  static std::unordered_map<std::string, std::shared_ptr<MacroArgRecord>> &
+  getMacroArgRecordMap() {
+    return MacroArgRecordMap;
+  }
   static std::map<std::string, std::shared_ptr<MacroExpansionRecord>> &
   getExpansionRangeToMacroRecord() {
     return ExpansionRangeToMacroRecord;
@@ -1568,6 +1588,11 @@ class DpctGlobalInfo {
   static std::map<std::string,
                   std::shared_ptr<DpctGlobalInfo::MacroExpansionRecord>>
       ExpansionRangeToMacroRecord;
+  // key: The hash string of the location of function-like macro argument
+  // value: Function-like macro argument information
+  static std::unordered_map<std::string,
+                            std::shared_ptr<DpctGlobalInfo::MacroArgRecord>>
+      MacroArgRecordMap;
   static std::map<std::string, SourceLocation> EndifLocationOfIfdef;
   static std::vector<std::pair<clang::tooling::UnifiedPath, size_t>>
       ConditionalCompilationLoc;
@@ -1894,7 +1919,7 @@ class MemVarInfo : public VarInfo {
   bool isUseHelperFunc() { return UseHelperFuncFlag; }
   void setUseDeviceGlobalFlag(bool Flag) { UseDeviceGlobalFlag = Flag; }
   bool isUseDeviceGlobal() { return UseDeviceGlobalFlag; }
-  void setInitForDeviceGlobal(std::string Init) { InitList = Init; }
+  void migrateToDeviceGlobal(const VarDecl *MemVar);
 
 private:
   bool isTreatPointerAsArray() {

clang/lib/DPCT/Utility.cpp

@@ -4811,16 +4811,6 @@ std::string getNameSpace(const NamespaceDecl *NSD) {
   return NameSpace;
 }
 
-std::string getInitForDeviceGlobal(const VarDecl *VD) {
-  auto Init = VD->getInit()->IgnoreImplicitAsWritten();
-  if (auto IL = dyn_cast<InitListExpr>(Init)) {
-    return dpct::ExprAnalysis::ref(IL);
-  } else if (dyn_cast<CXXConstructExpr>(Init)) {
-    return "";
-  }
-  return "{" + dpct::ExprAnalysis::ref(Init) + "}";
-}
-
 void getNameSpace(const NamespaceDecl *NSD,
                   std::vector<std::string> &Namespaces) {
   if (!NSD)

clang/lib/DPCT/Utility.h

@@ -568,7 +568,6 @@ bool containIterationSpaceBuiltinVar(const clang::Stmt *Node);
 bool containBuiltinWarpSize(const clang::Stmt *Node);
 bool isCapturedByLambda(const clang::TypeLoc *TL);
 std::string getNameSpace(const NamespaceDecl *NSD);
-std::string getInitForDeviceGlobal(const VarDecl *VD);
 void getNameSpace(const NamespaceDecl *NSD,
                   std::vector<std::string> &Namespaces);
 std::string getTemplateArgumentAsString(const clang::TemplateArgument &Arg,

clang/test/dpct/device_global/device_global.cu

@@ -60,9 +60,20 @@ __constant__ int arr_d[10] = {2};
 __constant__ A arr_e[10];
 __device__ B arr_f[10];
 
+// CHECK: #define TABLE_BEGIN(type, name, size) static const sycl::ext::oneapi::experimental::device_global<type[size]> name {
+// CHECK: #define TABLE_END() };
+#define TABLE_BEGIN(type, name, size) static const __device__ type name[size] = {
+#define TABLE_END() };
+
+// CHECK: TABLE_BEGIN(int, arr_g, 10)
+// CHECK: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
+// CHECK: TABLE_END()
+TABLE_BEGIN(int, arr_g, 10)
+0, 1, 2, 3, 4, 5, 6, 7, 8, 9
+TABLE_END()
 
 // CHECK: int device_func() {
-// CHECK:   float *p = arra_a.get();
+// CHECK:   float *p = arr_a.get();
 // CHECK:   arr_a[0] = 1;
 // CHECK:   return arr_a[0] + arr_b[0] + arr_c[0] + arr_d[0] + arr_e[0].data + arr_f[0].data;
 // CHECK: }
@@ -73,8 +84,9 @@ __device__ B arr_f[10];
 // CHECK:   *ptr = var_a.get() + var_b.get() + var_c.get() + var_d.get() + var_e.get().data + var_f.get().data + device_func();
 // CHECK: }
 __device__ int device_func() {
-  float *p = arra_a;
+  float *p = arr_a;
   arr_a[0] = 1;
+  arr_g[0];
   return arr_a[0] + arr_b[0] + arr_c[0] + arr_d[0] + arr_e[0].data + arr_f[0].data;
 }