From 7280da2f8888668b0ab2ed44fd5bfeab3c1717c8 Mon Sep 17 00:00:00 2001
From: Teja Alaghari <teja.alaghari@intel.com>
Date: Wed, 27 Nov 2024 11:05:36 +0800
Subject: [PATCH] Added support for more torch.cuda APIs

---
 .../case_006/case_006_torch_cuda.cpp          | 10 +++++
 .../python_migration/case_006/expected.py     | 14 ++++++
 .../dpct/python_migration/case_006/input.py   | 14 ++++++
 ...thon_build_script_migration_rule_ipex.yaml | 45 +++++++++++++++++++
 4 files changed, 83 insertions(+)
 create mode 100644 clang/test/dpct/python_migration/case_006/case_006_torch_cuda.cpp
 create mode 100644 clang/test/dpct/python_migration/case_006/expected.py
 create mode 100644 clang/test/dpct/python_migration/case_006/input.py

diff --git a/clang/test/dpct/python_migration/case_006/case_006_torch_cuda.cpp b/clang/test/dpct/python_migration/case_006/case_006_torch_cuda.cpp
new file mode 100644
index 000000000000..cbd0ddb9c747
--- /dev/null
+++ b/clang/test/dpct/python_migration/case_006/case_006_torch_cuda.cpp
@@ -0,0 +1,10 @@
+// RUN: rm -rf %T && mkdir -p %T
+// RUN: cd %T
+// RUN: cp %S/input.py ./input.py
+// RUN: dpct -in-root ./ -out-root out ./input.py --migrate-build-script-only
+// RUN: echo "begin" > %T/diff.txt
+// RUN: diff --strip-trailing-cr %S/expected.py %T/out/input.py >> %T/diff.txt
+// RUN: echo "end" >> %T/diff.txt
+
+// CHECK: begin
+// CHECK-NEXT: end
diff --git a/clang/test/dpct/python_migration/case_006/expected.py b/clang/test/dpct/python_migration/case_006/expected.py
new file mode 100644
index 000000000000..e49ba1436412
--- /dev/null
+++ b/clang/test/dpct/python_migration/case_006/expected.py
@@ -0,0 +1,14 @@
+from torch import xpu
+
+devs = torch.xpu.device_count()
+devs = xpu.device_count()
+
+d_cap = torch.xpu.get_device_capability()
+d_cap = xpu.get_device_capability()
+d0_cap = torch.xpu.get_device_capability(devs[0])
+d0_cap = xpu.get_device_capability(devs[0])
+
+arch_list = ['']
+arch_list = ['']
+
+cuda_ver = torch.version.xpu
diff --git a/clang/test/dpct/python_migration/case_006/input.py b/clang/test/dpct/python_migration/case_006/input.py
new file mode 100644
index 000000000000..341d4e49487f
--- /dev/null
+++ b/clang/test/dpct/python_migration/case_006/input.py
@@ -0,0 +1,14 @@
+from torch import cuda
+
+devs = torch.cuda.device_count()
+devs = cuda.device_count()
+
+d_cap = torch.cuda.get_device_capability()
+d_cap = cuda.get_device_capability()
+d0_cap = torch.cuda.get_device_capability(devs[0])
+d0_cap = cuda.get_device_capability(devs[0])
+
+arch_list = torch.cuda.get_arch_list()
+arch_list = cuda.get_arch_list()
+
+cuda_ver = torch.version.cuda
diff --git a/clang/tools/dpct/DpctOptRules/python_build_script_migration_rule_ipex.yaml b/clang/tools/dpct/DpctOptRules/python_build_script_migration_rule_ipex.yaml
index 00a86bb538d0..6ac3f486ffe5 100644
--- a/clang/tools/dpct/DpctOptRules/python_build_script_migration_rule_ipex.yaml
+++ b/clang/tools/dpct/DpctOptRules/python_build_script_migration_rule_ipex.yaml
@@ -101,6 +101,14 @@
   In: CUDAExtension
   Out: DPCPPExtension
 
+- Rule: rule_from_torch_import_cuda
+  Kind: PythonRule
+  Priority: Fallback
+  MatchMode: Partial
+  PythonSyntax: from_torch_import_cuda
+  In: from torch import cuda
+  Out: from torch import xpu
+
 - Rule: rule_cuda_is_available
   Kind: PythonRule
   Priority: Fallback
@@ -137,3 +145,40 @@
       MatchMode: Full
       In: ${arg}.cu
       Out: ${arg}.cpp
+
+- Rule: rule_cuda_device_count
+  Kind: PythonRule
+  Priority: Fallback
+  MatchMode: Partial
+  PythonSyntax: cuda_device_count
+  In: cuda.device_count
+  Out: xpu.device_count
+
+- Rule: rule_cuda_get_device_capability
+  Kind: PythonRule
+  Priority: Fallback
+  MatchMode: Partial
+  PythonSyntax: cuda_get_device_capability
+  In: cuda.get_device_capability
+  Out: xpu.get_device_capability
+
+- Rule: rule_cuda_get_arch_list
+  Kind: PythonRule
+  Priority: Fallback
+  MatchMode: Partial
+  PythonSyntax: cuda_get_arch_list
+  In: ${torch_prefix}cuda.get_arch_list()
+  Out: ${torch_prefix}['']
+  Subrules:
+    torch_prefix:
+      MatchMode: Full
+      In: torch.
+      Out: ""
+
+- Rule: rule_version_cuda
+  Kind: PythonRule
+  Priority: Fallback
+  MatchMode: Partial
+  PythonSyntax: version_cuda
+  In: version.cuda
+  Out: version.xpu