[SYCLomatic] Fix the FirstIncludeOffset (#2456)

zhiweij1 · web-flow · commit 491de119c3f7 · 2024-11-08T11:17:12.000+08:00
Signed-off-by: Jiang, Zhiwei &lt;zhiwei.jiang@intel.com&gt;
diff --git a/clang/lib/DPCT/ASTTraversal.cpp b/clang/lib/DPCT/ASTTraversal.cpp
@@ -9784,6 +9784,10 @@ void MemoryMigrationRule::mallocMigration(
       DpctGlobalInfo::addPriorityReplInfo(
           LocInfo.first.getCanonicalPath().str() + std::to_string(LocInfo.second), Info);
     }
+    // Insert header here since PriorityReplInfo delay the replacement addation
+    // to the post process. At that time, the MainFile is invalid.
+    DpctGlobalInfo::getInstance().insertHeader(C->getBeginLoc(),
+                                               HeaderType::HT_SYCL);
     instrumentAddressToSizeRecordForCodePin(C,0,1);
   } else if (Name == "cudaHostAlloc" || Name == "cudaMallocHost" ||
              Name == "cuMemHostAlloc" || Name == "cuMemAllocHost_v2" ||
@@ -9804,6 +9808,10 @@ void MemoryMigrationRule::mallocMigration(
                          EA.getSubExprRepl().end());
       DpctGlobalInfo::addPriorityReplInfo(
           LocInfo.first.getCanonicalPath().str() + std::to_string(LocInfo.second), Info);
+      // Insert header here since PriorityReplInfo delay the replacement
+      // addation to the post process. At that time, the MainFile is invalid.
+      DpctGlobalInfo::getInstance().insertHeader(C->getBeginLoc(),
+                                                 HeaderType::HT_SYCL);
     } else {
       ManagedPointerAnalysis MPA(C, IsAssigned);
       MPA.RecursiveAnalyze();
@@ -10636,6 +10644,10 @@ void MemoryMigrationRule::miscMigration(const MatchFinder::MatchResult &Result,
                          EA.getSubExprRepl().end());
       DpctGlobalInfo::addPriorityReplInfo(
           LocInfo.first.getCanonicalPath().str() + std::to_string(LocInfo.second), Info);
+      // Insert header here since PriorityReplInfo delay the replacement
+      // addation to the post process. At that time, the MainFile is invalid.
+      DpctGlobalInfo::getInstance().insertHeader(C->getBeginLoc(),
+                                                 HeaderType::HT_SYCL);
     } else {
       report(C->getBeginLoc(), Diagnostics::API_NOT_MIGRATED, false,
              MapNames::ITFName.at(Name));
diff --git a/clang/lib/DPCT/AnalysisInfo.cpp b/clang/lib/DPCT/AnalysisInfo.cpp
@@ -813,16 +813,18 @@ bool DpctFileInfo::isInAnalysisScope() {
   return DpctGlobalInfo::isInAnalysisScope(FilePath);
 }
 void DpctFileInfo::setFileEnterOffset(unsigned Offset) {
-  if (!HasInclusionDirective) {
-    FirstIncludeOffset = Offset;
+  auto MF = DpctGlobalInfo::getInstance().getMainFile();
+  if (!HasInclusionDirectiveSet.count(MF)) {
+    FirstIncludeOffset[MF] = Offset;
     LastIncludeOffset = Offset;
   }
 }
 void DpctFileInfo::setFirstIncludeOffset(unsigned Offset) {
-  if (!HasInclusionDirective) {
-    FirstIncludeOffset = Offset;
+  auto MF = DpctGlobalInfo::getInstance().getMainFile();
+  if (!HasInclusionDirectiveSet.count(MF)) {
+    FirstIncludeOffset[MF] = Offset;
     LastIncludeOffset = Offset;
-    HasInclusionDirective = true;
+    HasInclusionDirectiveSet.insert(MF);
   }
 }
 void DpctFileInfo::concatHeader(llvm::raw_string_ostream &OS) {}
@@ -853,9 +855,10 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
                                 ReplacementType IsForCodePin) {
   if (Type == HT_DPL_Algorithm || Type == HT_DPL_Execution || Type == HT_SYCL) {
     if (auto MF = DpctGlobalInfo::getInstance().getMainFile())
-      if (this != MF.get())
+      if (this != MF.get() && FirstIncludeOffset.count(MF)) {
         DpctGlobalInfo::getInstance().getMainFile()->insertHeader(
-            Type, FirstIncludeOffset);
+            Type, FirstIncludeOffset.at(MF));
+      }
   }
   if (HeaderInsertedBitMap[Type])
     return;
@@ -877,8 +880,11 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
   case HT_DPL_Algorithm:
   case HT_DPL_Execution:
     concatHeader(OS, getHeaderSpelling(Type));
-    return insertHeader(OS.str(), FirstIncludeOffset,
-                        InsertPosition::IP_AlwaysLeft);
+    if (auto Iter = FirstIncludeOffset.find(
+            DpctGlobalInfo::getInstance().getMainFile());
+        Iter != FirstIncludeOffset.end())
+      insertHeader(OS.str(), Iter->second, InsertPosition::IP_AlwaysLeft);
+    return;
   case HT_SYCL:
     // Add the label for profiling macro "DPCT_PROFILING_ENABLED", which will be
     // replaced by "#define DPCT_PROFILING_ENABLED" or not in the post
@@ -924,7 +930,10 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
            << DpctGlobalInfo::getGlobalQueueName() << ";" << getNL();
       }
     }
-    insertHeader(OS.str(), FirstIncludeOffset, InsertPosition::IP_Left);
+    if (auto Iter = FirstIncludeOffset.find(
+            DpctGlobalInfo::getInstance().getMainFile());
+        Iter != FirstIncludeOffset.end())
+      insertHeader(OS.str(), Iter->second, InsertPosition::IP_Left);
     if (!RTVersionValue.empty())
       MigratedMacroDefinitionOS << "#define DPCT_COMPAT_RT_VERSION "
                                 << RTVersionValue << getNL();
@@ -941,8 +950,11 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
                  InsertPosition::IP_AlwaysLeft);
     for (const auto &File :
          DpctGlobalInfo::getCustomHelperFunctionAddtionalIncludes()) {
-      insertHeader("#include \"" + File + +"\"" + getNL(), FirstIncludeOffset,
-                   InsertPosition::IP_Right);
+      if (auto Iter = FirstIncludeOffset.find(
+              DpctGlobalInfo::getInstance().getMainFile());
+          Iter != FirstIncludeOffset.end())
+        insertHeader("#include \"" + File + +"\"" + getNL(), Iter->second,
+                     InsertPosition::IP_Right);
     }
     return;
 
@@ -994,14 +1006,20 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
     }
     SchemaRelativePath += "codepin_autogen_util.hpp\"";
     concatHeader(OS, SchemaRelativePath);
-    return insertHeader(OS.str(), FirstIncludeOffset, InsertPosition::IP_Right,
-                        IsForCodePin);
+    if (auto Iter = FirstIncludeOffset.find(
+            DpctGlobalInfo::getInstance().getMainFile());
+        Iter != FirstIncludeOffset.end())
+      insertHeader(OS.str(), Iter->second, InsertPosition::IP_Right,
+                   IsForCodePin);
+    return;
   } break;
   default:
     break;
   }
 
-  if (Offset != FirstIncludeOffset)
+  if (FirstIncludeOffset.count(DpctGlobalInfo::getInstance().getMainFile()) &&
+      Offset !=
+          FirstIncludeOffset.at(DpctGlobalInfo::getInstance().getMainFile()))
     OS << getNL();
   concatHeader(OS, getHeaderSpelling(Type));
   return insertHeader(OS.str(), LastIncludeOffset, InsertPosition::IP_Right);
diff --git a/clang/lib/DPCT/AnalysisInfo.h b/clang/lib/DPCT/AnalysisInfo.h
@@ -586,10 +586,14 @@ class DpctFileInfo {
   clang::tooling::UnifiedPath FilePath;
   std::string FileContentCache;
 
-  unsigned FirstIncludeOffset = 0;
+  // Save the FirstIncludeOffset in each MainFile
+  std::map<std::shared_ptr<DpctFileInfo> /*MainFile*/, unsigned>
+      FirstIncludeOffset;
   unsigned LastIncludeOffset = 0;
   const unsigned FileBeginOffset = 0;
-  bool HasInclusionDirective = false;
+  // Save the status whether FirstIncludeOffset is set by setFirstIncludeOffset
+  // for each MainFile
+  std::set<std::shared_ptr<DpctFileInfo> /*MainFile*/> HasInclusionDirectiveSet;
   std::vector<std::string> InsertedHeaders;
   std::vector<std::string> InsertedHeadersCUDA;
   std::bitset<32> HeaderInsertedBitMap;
diff --git a/clang/lib/DPCT/RulesLangLib/CUBAPIMigration.cpp b/clang/lib/DPCT/RulesLangLib/CUBAPIMigration.cpp
@@ -603,12 +603,22 @@ void removeVarDecl(const VarDecl *VD) {
                 (new ReplaceStmt(DS, ""))->getReplacement(Context));
             DeclPRInfo->Priority = 1;
             DpctGlobalInfo::addPriorityReplInfo(Key, DeclPRInfo);
+            // Insert header here since PriorityReplInfo delay the replacement
+            // addation to the post process. At that time, the MainFile is
+            // invalid.
+            DpctGlobalInfo::getInstance().insertHeader(DS->getBeginLoc(),
+                                                       HeaderType::HT_SYCL);
           } else {
             auto SubDeclPRInfo = std::make_shared<PriorityReplInfo>();
             SubDeclPRInfo->Repls.emplace_back(
                 replaceText(Beg, End.getLocWithOffset(1), "", SM)
                     ->getReplacement(Context));
             DpctGlobalInfo::addPriorityReplInfo(Key, SubDeclPRInfo);
+            // Insert header here since PriorityReplInfo delay the replacement
+            // addation to the post process. At that time, the MainFile is
+            // invalid.
+            DpctGlobalInfo::getInstance().insertHeader(Beg,
+                                                       HeaderType::HT_SYCL);
           }
           break;
         }
@@ -666,6 +676,11 @@ void CubDeviceLevelRule::removeRedundantTempVar(const CallExpr *CE) {
         auto LocInfo = DpctGlobalInfo::getLocInfo((*Itr)->getBeginLoc());
         auto Info = std::make_shared<PriorityReplInfo>();
         Info->Priority = 1;
+        // Insert header here since PriorityReplInfo delay the replacement
+        // addation to the post process. At that time, the MainFile is
+        // invalid.
+        DpctGlobalInfo::getInstance().insertHeader((*Itr)->getBeginLoc(),
+                                                   HeaderType::HT_SYCL);
         if (IsUsed) {
           Info->Repls.emplace_back(ReplaceStmt(*Itr, "0").getReplacement(
               DpctGlobalInfo::getContext()));
diff --git a/clang/test/dpct/header_insert/b.cpp b/clang/test/dpct/header_insert/b.cpp
@@ -0,0 +1,4 @@
+// RUN: echo ""
+
+// CHECK: #include "b.h"
+#include "b.h"
diff --git a/clang/test/dpct/header_insert/b.cu b/clang/test/dpct/header_insert/b.cu
@@ -0,0 +1,38 @@
+// UNSUPPORTED: system-windows
+// RUN: cd %T
+// RUN: cp %S/b.cu .
+// RUN: cp %S/b.cpp .
+// RUN: cp %S/b.h .
+// RUN: echo "[" > compile_commands.json
+// RUN: echo "    {" >> compile_commands.json
+// RUN: echo "        \"command\": \"c++ b.cpp\"," >> compile_commands.json
+// RUN: echo "        \"directory\": \"%/T\"," >> compile_commands.json
+// RUN: echo "        \"file\": \"%/T/b.cpp\"" >> compile_commands.json
+// RUN: echo "    }," >> compile_commands.json
+// RUN: echo "    {" >> compile_commands.json
+// RUN: echo "        \"command\": \"nvcc b.cu\"," >> compile_commands.json
+// RUN: echo "        \"directory\": \"%/T\"," >> compile_commands.json
+// RUN: echo "        \"file\": \"%/T/b.cu\"" >> compile_commands.json
+// RUN: echo "    }" >> compile_commands.json
+// RUN: echo "]" >> compile_commands.json
+// RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
+// RUN: FileCheck %S/b.cu  --match-full-lines --input-file %T/out/b.dp.cpp
+// RUN: FileCheck %S/b.cpp --match-full-lines --input-file %T/out/b.cpp
+// RUN: FileCheck %S/b.h   --match-full-lines --input-file %T/out/b.h
+// RUN: rm -rf ./out
+
+// CHECK: #define CU_FILE
+// CHECK-NEXT: #include <sycl/sycl.hpp>
+// CHECK-NEXT: #include <dpct/dpct.hpp>
+// CHECK-NEXT: #include "b.h"
+// CHECK-EMPTY:
+// CHECK-NEXT: void foo() {
+// CHECK-NEXT:   sycl::float2 f2;
+// CHECK-NEXT: }
+
+#define CU_FILE
+#include "b.h"
+
+void foo() {
+  float2 f2;
+}
diff --git a/clang/test/dpct/header_insert/b.h b/clang/test/dpct/header_insert/b.h
@@ -0,0 +1,34 @@
+
+// CHECK: #ifndef CU_FILE
+// CHECK-NEXT: #include <cstdio>
+// CHECK-NEXT: #else
+// CHECK-NEXT: #include <sycl/sycl.hpp>
+// CHECK-NEXT: #include <dpct/dpct.hpp>
+// CHECK-NEXT: #include <iostream>
+// CHECK-NEXT: #endif
+// CHECK-EMPTY:
+// CHECK-NEXT: typedef struct dpct_type_{{[0-9]+}} {
+// CHECK-NEXT:   unsigned i;
+// CHECK-NEXT: } T1;
+// CHECK-EMPTY:
+// CHECK-NEXT: #ifdef CU_FILE
+// CHECK-NEXT: sycl::float2 ff;
+// CHECK-NEXT: #endif
+
+// This test targets to make sure SYCL header files are inserted into
+// right location (in the code section: CU_FILE is defined)
+// to avoid build fail for code path without CU_FILE.
+
+#ifndef CU_FILE
+#include <cstdio>
+#else
+#include <iostream>
+#endif
+
+typedef struct {
+  unsigned i;
+} T1;
+
+#ifdef CU_FILE
+float2 ff;
+#endif
diff --git a/clang/test/dpct/unused_atomic_macro.cu b/clang/test/dpct/unused_atomic_macro.cu
@@ -4,9 +4,7 @@
 // RUN: FileCheck --input-file %T/unused_atomic_macro/unused_atomic_macro.dp.cpp --match-full-lines %s
 // RUN: %if build_lit %{icpx -c -fsycl %T/unused_atomic_macro/unused_atomic_macro.dp.cpp -o %T/unused_atomic_macro/unused_atomic_macro.dp.o %}
 
-// CHECK: #include <sycl/sycl.hpp>
-// CHECK-NEXT: #include <dpct/dpct.hpp>
-// CHECK-NEXT:  /*
+// CHECK: /*
 // CHECK-NEXT: DPCT1058:{{[0-9]+}}: "atomicAdd" is not migrated because it is not called in the code.
 // CHECK-NEXT: */
 // CHECK-NEXT: #define ATOMIC_ADD(x, v) atomicAdd(&x, v);
