[SYCLomatic] Add more rules for pytorch API migration (#2740)

zhiweij1 · web-flow · commit 1a56ee6cf903 · 2025-04-01T13:53:13.000+08:00
Signed-off-by: Jiang, Zhiwei &lt;zhiwei.jiang@intel.com&gt;
diff --git a/clang/lib/DPCT/RulesInclude/InclusionHeaders.cpp b/clang/lib/DPCT/RulesInclude/InclusionHeaders.cpp
@@ -8,6 +8,8 @@
 
 #include "InclusionHeaders.h"
 #include "PreProcessor.h"
+#include "UserDefinedRules/UserDefinedRules.h"
+#include <optional>
 
 namespace clang {
 namespace dpct {
@@ -33,11 +35,10 @@ class LastInclusionLocationUpdater {
   bool UpdateNeeded;
 };
 
-std::string applyUserDefinedHeader(const std::string &FileName) {
-  // Apply user-defined rule if needed
+std::optional<std::pair<std::string, RulePriority>>
+getUserDefinedHeader(const std::string &FileName) {
   auto It = MapNames::HeaderRuleMap.find(FileName);
-  if (It != MapNames::HeaderRuleMap.end() &&
-      It->second.Priority == RulePriority::Takeover) {
+  if (It != MapNames::HeaderRuleMap.end()) {
     auto &Rule = It->second;
     std::string ReplHeaderStr = Rule.Prefix;
     llvm::raw_string_ostream OS(ReplHeaderStr);
@@ -54,11 +55,12 @@ std::string applyUserDefinedHeader(const std::string &FileName) {
     for (auto &Header : Rule.Includes) {
       PrintHeader(Header);
     }
-    PrintHeader(Rule.Out);
+    if (!Rule.Out.empty())
+      PrintHeader(Rule.Out);
     OS << Rule.Postfix;
-    return ReplHeaderStr;
+    return std::make_pair(ReplHeaderStr, Rule.Priority);
   }
-  return "";
+  return std::nullopt;
 }
 
 void insertHeaders(std::shared_ptr<DpctFileInfo> File,
@@ -150,6 +152,15 @@ void IncludesCallbacks::InclusionDirective(
     Updater.give_up();
   };
 
+  // Apply user-defined rule if needed
+  auto UserDefinedInfo = getUserDefinedHeader(FileName.str());
+  if (UserDefinedInfo.has_value()) {
+    if (UserDefinedInfo.value().second == RulePriority::Takeover) {
+      EmplaceReplacement(std::move(UserDefinedInfo.value().first));
+      return;
+    }
+  }
+
   if (Global.isInAnalysisScope(IncludedFile)) {
     IncludeFileMap[IncludedFile] = false;
     Global.getIncludingFileSet().insert(IncludedFile);
@@ -199,6 +210,7 @@ void IncludesCallbacks::InclusionDirective(
                         .getReplacement(DpctGlobalInfo::getContext());
         DpctGlobalInfo::getIncludeMapSet().push_back({IncludedFile, Repl});
       }
+      UserDefinedInfo.reset();
     }
     if (Global.isInRoot(IncludedFile))
       return;
@@ -208,11 +220,8 @@ void IncludesCallbacks::InclusionDirective(
       !Global.getSourceManager().isWrittenInMainFile(HashLoc))
     return;
 
-
-  // Apply user-defined rule if needed
-  if (auto ReplacedStr = applyUserDefinedHeader(FileName.str());
-      !ReplacedStr.empty()) {
-    EmplaceReplacement(std::move(ReplacedStr));
+  if (UserDefinedInfo.has_value()) {
+    EmplaceReplacement(std::move(UserDefinedInfo.value().first));
     return;
   }
 
diff --git a/clang/test/dpct/pytorch/ATen.cu b/clang/test/dpct/pytorch/ATen.cu
@@ -4,7 +4,7 @@
 // RUN: cp -r %S/pytorch_inc %T/pytorch/ATen/
 // RUN: cd %T/pytorch/ATen
 // RUN: mkdir dpct_out
-// RUN: dpct --out-root dpct_out %T/pytorch/ATen/src/ATen.cu --extra-arg="-I%T/pytorch/ATen/pytorch_inc" --cuda-include-path="%cuda-path/include" --rule-file=%S/../../../tools/dpct/extensions/pytorch_api_rules/pytorch_api.yaml --analysis-scope-path %T/pytorch/ATen/pytorch_inc --analysis-scope-path %T/pytorch/ATen/src --in-root %T/pytorch/ATen/src
+// RUN: dpct --format-range=none --out-root dpct_out %T/pytorch/ATen/src/ATen.cu --extra-arg="-I%T/pytorch/ATen/pytorch_inc" --cuda-include-path="%cuda-path/include" --rule-file=%S/../../../tools/dpct/extensions/pytorch_api_rules/pytorch_api.yaml --analysis-scope-path %T/pytorch/ATen/pytorch_inc --analysis-scope-path %T/pytorch/ATen/src --in-root %T/pytorch/ATen/src
 // RUN: FileCheck --input-file %T/pytorch/ATen/dpct_out/ATen.dp.cpp --match-full-lines %T/pytorch/ATen/src/ATen.cu
 
 // CHECK: #include <c10/xpu/XPUStream.h>
@@ -18,6 +18,18 @@
 // CHECK-NEXT: #include <c10/util/Half.h>
 #include <ATen/cuda/CUDATensorMethods.cuh>
 
+// CHECK: // BEGIN_1
+// CHECK-EMPTY:
+// CHECK-EMPTY:
+// CHECK-NEXT: // END_1
+// BEGIN_1
+#include <ATen/cuda/Exceptions.h>
+#include <THC/THCAtomics.cuh>
+// END_1
+
+// CHECK: #include <c10/xpu/XPUMacros.h>
+#include <c10/cuda/CUDAMacros.h>
+
 #define AT_CUDA_CHECK(stmt)  (stmt)
 
 // CHECK: #define BE_AT_CHECK
@@ -31,20 +43,19 @@ void test_CUDAStream_as_arg() {
   dim3 blockSize(8, 8, 1);
   void *args[] = {nullptr}; 
 
-  // CHECK: ([&]() {
-  // CHECK-NEXT:   ((sycl::queue *)(c10::xpu::getCurrentXPUStream()))
-  // CHECK-NEXT:       ->parallel_for(sycl::nd_range<3>(gridSize * blockSize, blockSize),
-  // CHECK-NEXT:                      [=](sycl::nd_item<3> item_ct1) {
-  // CHECK-NEXT:                        kernel();
-  // CHECK-NEXT:                      });
+  // CHECK: ([&](){
+  // CHECK-NEXT:   ((sycl::queue*)(c10::xpu::getCurrentXPUStream()))->parallel_for(
+  // CHECK-NEXT:     sycl::nd_range<3>(gridSize * blockSize, blockSize), 
+  // CHECK-NEXT:     [=](sycl::nd_item<3> item_ct1) {
+  // CHECK-NEXT:       kernel();
+  // CHECK-NEXT:     });
   // CHECK-NEXT:   return 0;
   // CHECK-NEXT: }());
   AT_CUDA_CHECK(cudaLaunchKernel((const void *)kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));
 }
 
 int main() {
-  // CHECK: dpct::queue_ptr st =
-  // CHECK-NEXT: &static_cast<sycl::queue &>(c10::xpu::getCurrentXPUStream());
+  // CHECK: dpct::queue_ptr st = &static_cast<sycl::queue&>(c10::xpu::getCurrentXPUStream());
   cudaStream_t st = 0;
 
   // stream APIs
@@ -55,14 +66,18 @@ int main() {
   // CHECK: auto deviceStream = c10::xpu::getCurrentXPUStream(devInd);
   auto deviceStream = at::cuda::getCurrentCUDAStream(devInd);
 
-  // CHECK: dpct::queue_ptr curr_cuda_st =
-  // CHECK-NEXT:    &static_cast<sycl::queue &>(c10::xpu::getCurrentXPUStream().queue());
+  // CHECK: dpct::queue_ptr curr_cuda_st = &static_cast<sycl::queue &>(c10::xpu::getCurrentXPUStream(). queue());
   cudaStream_t curr_cuda_st = at::cuda::getCurrentCUDAStream().stream();
-  // CHECK: dpct::queue_ptr dev_cuda_st = &static_cast<sycl::queue &>(
-  // CHECK-NEXT:    c10::xpu::getCurrentXPUStream(devInd).queue());
+  // CHECK: dpct::queue_ptr dev_cuda_st = &static_cast<sycl::queue &>(c10::xpu::getCurrentXPUStream(devInd). queue());
   cudaStream_t dev_cuda_st = at::cuda::getCurrentCUDAStream(devInd).stream();
 
   test_CUDAStream_as_arg();
 
   return 0;
 }
+
+// CHECK: void foo2(c10::DeviceGuard device_guard, float *f)  try {
+// CHECK-NEXT: (DPCT_CHECK_ERROR(f = (float *)sycl::malloc_device(4, static_cast<sycl::queue&>(c10::xpu::getCurrentXPUStream()))));
+void foo2(at::cuda::CUDAGuard device_guard, float *f) {
+  C10_CUDA_CHECK(cudaMalloc(&f, 4));
+}
diff --git a/clang/test/dpct/pytorch/pytorch_inc/ATen/Tensor.h b/clang/test/dpct/pytorch/pytorch_inc/ATen/Tensor.h
@@ -2,6 +2,7 @@
 namespace at {
 class Tensor {
 public:
-  bool is_cuda();
+  int get_device() const { return 0; }
+  bool is_cuda() const  { return true; };
 };
 } // namespace at
diff --git a/clang/test/dpct/pytorch/pytorch_inc/ATen/cuda/CUDAContext.h b/clang/test/dpct/pytorch/pytorch_inc/ATen/cuda/CUDAContext.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAGuard.h>
 
 namespace at {
 using namespace c10;
diff --git a/clang/test/dpct/pytorch/pytorch_inc/ATen/cuda/Exceptions.h b/clang/test/dpct/pytorch/pytorch_inc/ATen/cuda/Exceptions.h
@@ -0,0 +1,2 @@
+#pragma once
+
diff --git a/clang/test/dpct/pytorch/pytorch_inc/THC/THCAtomics.cuh b/clang/test/dpct/pytorch/pytorch_inc/THC/THCAtomics.cuh
@@ -0,0 +1 @@
+// RUN: echo "empty command"
diff --git a/clang/test/dpct/pytorch/pytorch_inc/c10/cuda/CUDAGuard.h b/clang/test/dpct/pytorch/pytorch_inc/c10/cuda/CUDAGuard.h
@@ -9,6 +9,7 @@ class optional {
 } // namespace std
 
 namespace c10 {
+using DeviceIndex = int8_t;
 class Device {
 public:
   Device(std::string str) {}
@@ -19,5 +20,15 @@ class OptionalCUDAGuard {
 public:
   OptionalCUDAGuard(std::optional<c10::Device> device) {}
 };
+struct CUDAGuard {
+  explicit CUDAGuard() = delete;
+  explicit CUDAGuard(DeviceIndex device_index) {}
+  explicit CUDAGuard(Device device) {}
+  CUDAGuard(const CUDAGuard&) = delete;
+  CUDAGuard& operator=(const CUDAGuard&) = delete;
+  CUDAGuard(CUDAGuard&& other) = delete;
+  CUDAGuard& operator=(CUDAGuard&& other) = delete;
+  ~CUDAGuard() = default;
+};
 } // namespace cuda
 } // namespace c10
diff --git a/clang/test/dpct/pytorch/pytorch_inc/c10/cuda/CUDAMacros.h b/clang/test/dpct/pytorch/pytorch_inc/c10/cuda/CUDAMacros.h
@@ -3,3 +3,7 @@
 #define C10_CUDA_IMPORT
 #define C10_CUDA_API
 #define C10_CUDA_BUILD_MAIN_LIB
+#define C10_CUDA_CHECK(EXPR)                                        \
+  do {                                                              \
+    const cudaError_t __err = EXPR;                                 \
+  } while (0)
diff --git a/clang/tools/dpct/extensions/pytorch_api_rules/pytorch_api.yaml b/clang/tools/dpct/extensions/pytorch_api_rules/pytorch_api.yaml
@@ -168,3 +168,28 @@
   In: get_in_order_queue
   Out: static_cast<sycl::queue&>(c10::xpu::getCurrentXPUStream())
   Includes: [<c10/xpu/XPUStream.h>]
+
+- Rule: rule_THC_THCAtomics_cuh
+  Kind: Header
+  Priority: Takeover
+  In: THC/THCAtomics.cuh
+  Out: |
+
+- Rule: rule_ATen_cuda_Exceptions_h
+  Kind: Header
+  Priority: Takeover
+  In: ATen/cuda/Exceptions.h
+  Out: |
+
+- Rule: rule_remove_C10_CUDA_CHECK
+  Kind: Macro
+  Priority: Takeover
+  In: C10_CUDA_CHECK
+  Out: |
+
+- Rule: rule_at_cuda_CUDAGuard
+  Kind: Type
+  Priority: Takeover
+  In: c10::cuda::CUDAGuard
+  Out: c10::DeviceGuard
+  Includes: [<c10/core/DeviceGuard.h>]
