[SYCLomatic] Enhance the migration of __shared__ memory

clang/lib/DPCT/AnalysisInfo.cpp

@@ -2811,6 +2811,7 @@ void CtTypeInfo::setArrayInfo(const DependentSizedArrayTypeLoc &TL,
                               bool NeedSizeFold) {
   ContainSizeofType = containSizeOfType(TL.getSizeExpr());
   ExprAnalysis EA;
+  //
   EA.analyze(TL.getSizeExpr());
   auto TDSI = EA.getTemplateDependentStringInfo();
   if (TDSI->containsTemplateDependentMacro())

clang/lib/DPCT/RuleInfra/ExprAnalysis.cpp

@@ -421,6 +421,74 @@ void StringReplacements::replaceString() {
   ReplMap.clear();
 }
 
+void StringReplacements::adjustSourceStrAndTDRs() {
+  std::vector<size_t> Occurrences;
+  size_t Pos = 0;
+  while ((Pos = SourceStr.find(ConstExprExpansionInfoVDStr, Pos)) !=
+         std::string::npos) {
+    Occurrences.push_back(Pos);
+    Pos += ConstExprExpansionInfoVDStr.length();
+  }
+
+  int LengthDiff = ConstExprExpansionInfoVDInitStr.length() -
+                   ConstExprExpansionInfoVDStr.length();
+
+  for (auto Occ : Occurrences) {
+    const auto &Iter = TDRs.find(Occ);
+    if (Iter != TDRs.end()) {
+      TDRs.erase(Iter);
+    }
+  }
+  std::map<size_t, std::shared_ptr<TemplateDependentReplacement>> NewTDRs;
+  for (const auto &[Pos, Value] : TDRs) {
+    int NewPos = Pos;
+    for (auto Occ : Occurrences) {
+      if (Occ < static_cast<size_t>(Pos)) {
+        NewPos += LengthDiff;
+      } else {
+        break;
+      }
+    }
+    auto NewValue = std::make_shared<TemplateDependentReplacement>(
+        SourceStr, NewPos, Value->getLength(), Value->getTemplateIndex());
+    NewTDRs[NewPos] = NewValue;
+  }
+  TDRs = std::move(NewTDRs);
+
+  std::map<size_t, std::shared_ptr<TemplateDependentReplacement>>
+      NewTDRsInConstExprExpansion;
+  for (const auto &[Pos, Value] : TDRsInConstExprExpansion) {
+    for (auto Occ : Occurrences) {
+      int NewPos = Occ + Pos;
+      for (auto PrevOcc : Occurrences) {
+        if (PrevOcc < static_cast<size_t>(Occ)) {
+          NewPos += LengthDiff;
+        } else {
+          break;
+        }
+      }
+      auto NewValue = std::make_shared<TemplateDependentReplacement>(
+          SourceStr, NewPos, Value->getLength(), Value->getTemplateIndex());
+      NewTDRsInConstExprExpansion[NewPos] = NewValue;
+      break;
+    }
+  }
+  TDRsInConstExprExpansion = std::move(NewTDRsInConstExprExpansion);
+
+  Pos = 0;
+  while ((Pos = SourceStr.find(ConstExprExpansionInfoVDStr, Pos)) !=
+         std::string::npos) {
+    SourceStr.replace(Pos, ConstExprExpansionInfoVDStr.length(),
+                      ConstExprExpansionInfoVDInitStr);
+    Pos += ConstExprExpansionInfoVDInitStr.length();
+  }
+
+  TDRs.insert(TDRsInConstExprExpansion.begin(), TDRsInConstExprExpansion.end());
+  for (const auto &[Pos, Value] : TDRs) {
+    Value->alterSource(SourceStr);
+  }
+}
+
 ExprAnalysis::ExprAnalysis(const Expr *Expression)
     : Context(DpctGlobalInfo::getContext()),
       SM(DpctGlobalInfo::getSourceManager()) {
@@ -553,7 +621,18 @@ void ExprAnalysis::analyzeExpr(const DeclRefExpr *DRE) {
   }
   if (auto TemplateDecl = dyn_cast<NonTypeTemplateParmDecl>(DRE->getDecl()))
     addReplacement(DRE, TemplateDecl->getIndex());
-  else if (auto ECD = dyn_cast<EnumConstantDecl>(DRE->getDecl())) {
+  else if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl());
+           VD && VD->isConstexpr()) {
+    if (VD->getInit() && VD->getInit()->getBeginLoc().isValid() &&
+        !ConstExprExpansion) {
+      std::string VDInitStr = ExprAnalysis::ref(VD->getInit());
+      std::string VDStr = VD->getNameAsString();
+      ReplSet.addConstExprExpansionInfo(VDStr, VDInitStr);
+      ConstExprExpansion = true;
+      dispatch(VD->getInit());
+      ConstExprExpansion = false;
+    }
+  } else if (auto ECD = dyn_cast<EnumConstantDecl>(DRE->getDecl())) {
     std::unordered_set<std::string> targetStr = {
         "thread_scope_system",  "thread_scope_device",  "thread_scope_block",
         "memory_order_relaxed", "memory_order_acq_rel", "memory_order_release",

clang/lib/DPCT/RuleInfra/ExprAnalysis.h

@@ -77,6 +77,7 @@ class TemplateDependentReplacement {
   }
   inline size_t getOffset() const { return Offset; }
   inline size_t getLength() const { return Length; }
+  inline unsigned getTemplateIndex() const { return TemplateIndex; }
   const TemplateArgumentInfo &
   getTargetArgument(const std::vector<TemplateArgumentInfo> &TemplateList);
   void replace(const std::vector<TemplateArgumentInfo> &TemplateList);
@@ -136,6 +137,20 @@ class StringReplacements {
         std::make_pair(Offset, std::make_shared<TemplateDependentReplacement>(
                                    SourceStr, Offset, Length, TemplateIndex)));
   }
+
+  inline void addTemplateDependentReplacementInConstExprExpansion(
+      size_t Offset, size_t Length, unsigned TemplateIndex) {
+    TDRsInConstExprExpansion.insert(
+        std::make_pair(Offset, std::make_shared<TemplateDependentReplacement>(
+                                   SourceStr, Offset, Length, TemplateIndex)));
+  }
+
+  inline void addConstExprExpansionInfo(std::string VDStr,
+                                        std::string VDInitStr) {
+    ConstExprExpansionInfoVDStr = VDStr;
+    ConstExprExpansionInfoVDInitStr = VDInitStr;
+  }
+
   // Add a string replacement
   void addStringReplacement(size_t Offset, size_t Length, std::string Text) {
     auto Result = ReplMap.insert(std::make_pair(
@@ -155,6 +170,11 @@ class StringReplacements {
   // Generate replacement text info which dependent on template args.
   std::shared_ptr<TemplateDependentStringInfo>
   getTemplateDependentStringInfo() {
+    if (!ConstExprExpansionInfoVDStr.empty() &&
+        !ConstExprExpansionInfoVDInitStr.empty() &&
+        !TDRsInConstExprExpansion.empty()) {
+      adjustSourceStrAndTDRs();
+    }
     replaceString();
     return std::make_shared<TemplateDependentStringInfo>(SourceStr, TDRs);
   }
@@ -170,11 +190,16 @@ class StringReplacements {
   StringReplacements &operator=(StringReplacements) = delete;
 
   void replaceString();
+  void adjustSourceStrAndTDRs();
 
   unsigned ShiftLength;
   std::string SourceStr;
   std::map<size_t, std::shared_ptr<StringReplacement>> ReplMap;
   std::map<size_t, std::shared_ptr<TemplateDependentReplacement>> TDRs;
+  std::map<size_t, std::shared_ptr<TemplateDependentReplacement>>
+      TDRsInConstExprExpansion;
+  std::string ConstExprExpansionInfoVDStr;
+  std::string ConstExprExpansionInfoVDInitStr;
 };
 
 /// Analyze expression and generate its migrated string
@@ -583,7 +608,11 @@ class ExprAnalysis {
 
   inline void addReplacement(size_t Offset, size_t Length,
                              unsigned TemplateIndex) {
-    ReplSet.addTemplateDependentReplacement(Offset, Length, TemplateIndex);
+    if (ConstExprExpansion)
+      ReplSet.addTemplateDependentReplacementInConstExprExpansion(
+          Offset, Length, TemplateIndex);
+    else
+      ReplSet.addTemplateDependentReplacement(Offset, Length, TemplateIndex);
   }
 
   // Analyze the expression, jump to corresponding analysis function according
@@ -691,6 +720,7 @@ class ExprAnalysis {
   std::string RewritePrefix;
   std::string RewritePostfix;
   std::set<HelperFeatureEnum> HelperFeatureSet;
+  bool ConstExprExpansion = false;
 };
 
 // Analyze pointer allocated by cudaMallocManaged.

clang/test/dpct/sharedmem_var_static.cu

@@ -4,11 +4,12 @@
 // RUN: dpct --format-range=none --usm-level=none -out-root %T/sharedmem_var_static %s --cuda-include-path="%cuda-path/include" --sycl-named-lambda -- -x cuda --cuda-host-only
 // RUN: FileCheck %s --match-full-lines --input-file %T/sharedmem_var_static/sharedmem_var_static.dp.cpp
 // RUN: %if build_lit %{icpx -c -fsycl -DNO_BUILD_TEST  %T/sharedmem_var_static/sharedmem_var_static.dp.cpp -o %T/sharedmem_var_static/sharedmem_var_static.dp.o %}
-#ifndef NO_BUILD_TEST
+
 #include <stdio.h>
 #include <complex>
 #define SIZE 64
 
+#ifndef NO_BUILD_TEST
 class TestObject{
 public:
   // CHECK: static void run(int *in, int *out, const sycl::nd_item<3> &item_ct1, int &a0) {
@@ -226,3 +227,15 @@ void fooh() {
   fook<SZ><<<1, 1>>>();
 }
 #endif
+
+constexpr int kWarpSize = 32;
+
+template <int ThreadsPerBlock, int NumWarpQ> __global__ void kerfunc() {
+  constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
+  __shared__ int smem[kNumWarps * NumWarpQ];
+}
+
+void foo2() {
+  // CHECK: sycl::local_accessor<int, 1> smem_acc_ct1(sycl::range<1>(ThreadsPerBlock / kWarpSize * 8), cgh);
+  kerfunc<128, 8><<<32, 32>>>();
+}