[SYCLomatic] In a translation unit, if the header file <bits/stdc++.h> is included before SYCL header file, then insert the SYCL header file at the beginning of the main source file of the translation unit. (#2790)

Signed-off-by: Jiang, Zhiwei <[email protected]>

Author: zhiweij1

clang/lib/DPCT/AnalysisInfo.cpp

@@ -878,8 +878,14 @@ void DpctFileInfo::insertHeader(HeaderType Type, unsigned Offset,
   if (Type == HT_DPL_Algorithm || Type == HT_DPL_Execution || Type == HT_SYCL) {
     if (auto MF = DpctGlobalInfo::getInstance().getMainFile())
       if (this != MF.get() && FirstIncludeOffset.count(MF)) {
-        DpctGlobalInfo::getInstance().getMainFile()->insertHeader(
-            Type, FirstIncludeOffset.at(MF));
+        // If <bits/stdc++.h> is included before <sycl/sycl.hpp>, the
+        // compilation will fail.
+        auto Iter =
+            DpctGlobalInfo::getAfterBitsStdcxxFilesMap().find(MF->FilePath);
+        if (Iter != DpctGlobalInfo::getAfterBitsStdcxxFilesMap().end()) {
+          if (Iter->second.count(this->FilePath))
+            MF->insertHeader(Type, FirstIncludeOffset.at(MF));
+        }
       }
   }
   if (DpctGlobalInfo::getHeaderInsertedBitMap()[FilePath][Type])
@@ -2547,6 +2553,10 @@ std::unordered_set<std::string>
 std::unordered_map<clang::tooling::UnifiedPath,
                    std::bitset<HeaderType::NUM_HEADERS>>
     DpctGlobalInfo::HeaderInsertedBitMap = {};
+bool DpctGlobalInfo::IsAfterBitsStdcxx = false;
+std::map<clang::tooling::UnifiedPath /*MainFile*/,
+         std::set<clang::tooling::UnifiedPath>>
+    DpctGlobalInfo::AfterBitsStdcxxFiles;
 ///// class DpctNameGenerator /////
 void DpctNameGenerator::printName(const FunctionDecl *FD,
                                   llvm::raw_ostream &OS) {

clang/lib/DPCT/AnalysisInfo.h

@@ -1398,6 +1398,12 @@ class DpctGlobalInfo {
   getHeaderInsertedBitMap() {
     return HeaderInsertedBitMap;
   }
+  static bool &getIsAfterBitsStdcxxStatus() { return IsAfterBitsStdcxx; }
+  static std::map<clang::tooling::UnifiedPath,
+                  std::set<clang::tooling::UnifiedPath>> &
+  getAfterBitsStdcxxFilesMap() {
+    return AfterBitsStdcxxFiles;
+  }
   std::shared_ptr<DpctFileInfo>
   insertFile(const clang::tooling::UnifiedPath &FilePath) {
     return insertObject(FileMap, FilePath);
@@ -1407,7 +1413,10 @@ class DpctGlobalInfo {
     return findObject(FileMap, FilePath);
   }
   std::shared_ptr<DpctFileInfo> getMainFile() const { return MainFile; }
-  void setMainFile(std::shared_ptr<DpctFileInfo> Main) { MainFile = Main; }
+  void setMainFile(std::shared_ptr<DpctFileInfo> Main) {
+    getIsAfterBitsStdcxxStatus() = false;
+    MainFile = Main;
+  }
   void recordIncludingRelationship(
       const clang::tooling::UnifiedPath &CurrentFileName,
       const clang::tooling::UnifiedPath &IncludedFileName);
@@ -1726,6 +1735,17 @@ class DpctGlobalInfo {
   static std::unordered_map<clang::tooling::UnifiedPath,
                             std::bitset<HeaderType::NUM_HEADERS>>
       HeaderInsertedBitMap;
+  // `IsAfterBitsStdcxx` is used as a flag. It is set to true when PP meet
+  // <bits/stdc++.h>. It is reset to false when starting to process a new
+  // translation unit.
+  static bool IsAfterBitsStdcxx;
+  // If `IsAfterBitsStdcxx` is true, it means <bits/stdc++.h> is already
+  // included, so the tool will record all files included since this time point.
+  // When inserting `sycl.hpp`, the tool can check if the insert location is in
+  // the map, if it is, the tool will also insert `sycl.hpp` at the main file.
+  static std::map<clang::tooling::UnifiedPath /*MainFile*/,
+                  std::set<clang::tooling::UnifiedPath>>
+      AfterBitsStdcxxFiles;
 };
 
 /// Generate mangle name of FunctionDecl as key of DeviceFunctionInfo.

clang/lib/DPCT/PreProcessor.cpp

@@ -825,6 +825,11 @@ void IncludesCallbacks::FileChanged(SourceLocation Loc, FileChangeReason Reason,
     }
 
     clang::tooling::UnifiedPath InFile = SM.getFilename(Loc).str();
+    if (DpctGlobalInfo::getIsAfterBitsStdcxxStatus()) {
+      DpctGlobalInfo::getAfterBitsStdcxxFilesMap()
+          [DpctGlobalInfo::getInstance().getMainFile()->getFilePath()]
+              .insert(InFile);
+    }
     if (IsFileInCmd || ProcessAll ||
         GetSourceFileType(InFile.getCanonicalPath()) & SPT_CudaSource) {
       IncludeFileMap[DpctGlobalInfo::removeSymlinks(

clang/lib/DPCT/RulesInclude/InclusionHeaders.cpp

@@ -132,6 +132,10 @@ void IncludesCallbacks::InclusionDirective(
   FileInfo->setFirstIncludeOffset(LocInfo.second);
   LastInclusionLocationUpdater Updater(FileInfo, FilenameRange.getEnd());
 
+  if (FileName == "bits/stdc++.h") {
+    DpctGlobalInfo::getInstance().getIsAfterBitsStdcxxStatus() = true;
+  }
+
   clang::tooling::UnifiedPath IncludedFile;
   if (auto OptionalAbs = Global.getAbsolutePath(*File))
     IncludedFile = OptionalAbs.value();

clang/test/dpct/d_dh_constant_db.cpp

@@ -17,7 +17,6 @@
 // RUN: echo "        \"file\": \"%/T/d_dh_constant_db/d_dh_constant_db.cpp\"" >> compile_commands.json
 // RUN: echo "    }" >> compile_commands.json
 // RUN: echo "]" >> compile_commands.json
-// RUN: rm %T/d_dh_constant_db/out/MainSourceFiles.yaml
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: FileCheck %s --match-full-lines --input-file %T/d_dh_constant_db/out/constant_header.h
 // RUN: cd ..

clang/test/dpct/d_dh_constant_db_win.cpp

@@ -17,7 +17,6 @@
 // RUN: echo "        \"file\": \"%/T/d_dh_constant_db_win/d_dh_constant_db_win.cpp\"" >> compile_commands.json
 // RUN: echo "    }" >> compile_commands.json
 // RUN: echo "]" >> compile_commands.json
-// RUN: rm %T/d_dh_constant_db_win/out/MainSourceFiles.yaml
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: FileCheck %s --match-full-lines --input-file %T/d_dh_constant_db_win/out/constant_header.h
 // RUN: cd ..

clang/test/dpct/d_hd_constant_db.cpp

@@ -17,7 +17,6 @@
 // RUN: echo "        \"file\": \"%/T/d_hd_constant_db/d_hd_constant_db.cpp\"" >> compile_commands.json
 // RUN: echo "    }" >> compile_commands.json
 // RUN: echo "]" >> compile_commands.json
-// RUN: rm %T/d_hd_constant_db/out/MainSourceFiles.yaml
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: FileCheck %s --match-full-lines --input-file %T/d_hd_constant_db/out/constant_header.h
 // RUN: cd ..

clang/test/dpct/d_hd_constant_db_win.cpp

@@ -17,7 +17,6 @@
 // RUN: echo "        \"file\": \"%/T/d_hd_constant_db_win/d_hd_constant_db_win.cpp\"" >> compile_commands.json
 // RUN: echo "    }" >> compile_commands.json
 // RUN: echo "]" >> compile_commands.json
-// RUN: rm %T/d_hd_constant_db_win/out/MainSourceFiles.yaml
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: FileCheck %s --match-full-lines --input-file %T/d_hd_constant_db_win/out/constant_header.h
 // RUN: cd ..

clang/test/dpct/dh_d_constant_db.cpp

@@ -18,7 +18,6 @@
 // RUN: echo "]" >> compile_commands.json
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: rm compile_commands.json
-// RUN: rm %T/dh_d_constant_db/out/MainSourceFiles.yaml
 // RUN: dpct dh_d_constant_db.cpp --out-root=./out --cuda-include-path="%cuda-path/include" -- -x cuda --cuda-host-only
 // RUN: FileCheck %s --match-full-lines --input-file %T/dh_d_constant_db/out/constant_header.h
 // RUN: cd ..

clang/test/dpct/dh_d_constant_db_win.cpp

@@ -18,7 +18,6 @@
 // RUN: echo "]" >> compile_commands.json
 // RUN: dpct -p=. --out-root=./out --cuda-include-path="%cuda-path/include"
 // RUN: rm compile_commands.json
-// RUN: rm %T/dh_d_constant_db_win/out/MainSourceFiles.yaml
 // RUN: dpct dh_d_constant_db_win.cpp --out-root=./out --cuda-include-path="%cuda-path/include" -- -x cuda --cuda-host-only
 // RUN: FileCheck %s --match-full-lines --input-file %T/dh_d_constant_db_win/out/constant_header.h
 // RUN: cd ..