From e238351101bcd23a9a0a32dba0b071685d670d6c Mon Sep 17 00:00:00 2001
From: Logan Adams <114770087+loadams@users.noreply.github.com>
Date: Fri, 20 Oct 2023 12:05:54 -0700
Subject: [PATCH] ROCm 6.0 prep changes (#4537)

* ROCm 6.0 prep changes

* PR feedback

* Try updating apex
---
 .github/workflows/amd-mi200.yml | 1 +
 op_builder/cpu_adagrad.py       | 6 +-----
 op_builder/cpu_adam.py          | 6 +-----
 op_builder/random_ltd.py        | 3 ---
 op_builder/transformer.py       | 3 ---
 5 files changed, 3 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/amd-mi200.yml b/.github/workflows/amd-mi200.yml
index 77f33f744ea8..8c4292d4675c 100644
--- a/.github/workflows/amd-mi200.yml
+++ b/.github/workflows/amd-mi200.yml
@@ -45,6 +45,7 @@ jobs:
         run: |
           git clone https://github.com/ROCmSoftwarePlatform/apex.git
           cd apex
+          git checkout torch_2.1_higher
           CURRENT_VER=$(git rev-parse HEAD)
           INSTALLED_VER=$(cat /blob/amd-apex/.venv_installed_version)
           if [[ "$CURRENT_VER" != "$INSTALLED_VER" ]]; then
diff --git a/op_builder/cpu_adagrad.py b/op_builder/cpu_adagrad.py
index 6d70c93faac2..d3f163f7464a 100644
--- a/op_builder/cpu_adagrad.py
+++ b/op_builder/cpu_adagrad.py
@@ -39,9 +39,5 @@ def include_paths(self):
         elif not self.is_rocm_pytorch():
             CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
         else:
-            CUDA_INCLUDE = [
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"),
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"),
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"),
-            ]
+            CUDA_INCLUDE = []
         return ['csrc/includes'] + CUDA_INCLUDE
diff --git a/op_builder/cpu_adam.py b/op_builder/cpu_adam.py
index e500a1eea907..7c34c4ce43a1 100644
--- a/op_builder/cpu_adam.py
+++ b/op_builder/cpu_adam.py
@@ -40,9 +40,5 @@ def include_paths(self):
         elif not self.is_rocm_pytorch():
             CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
         else:
-            CUDA_INCLUDE = [
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"),
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"),
-                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"),
-            ]
+            CUDA_INCLUDE = []
         return ['csrc/includes'] + CUDA_INCLUDE
diff --git a/op_builder/random_ltd.py b/op_builder/random_ltd.py
index 3fdc777215da..54af7150fb36 100644
--- a/op_builder/random_ltd.py
+++ b/op_builder/random_ltd.py
@@ -31,7 +31,4 @@ def sources(self):
 
     def include_paths(self):
         includes = ['csrc/includes']
-        if self.is_rocm_pytorch():
-            from torch.utils.cpp_extension import ROCM_HOME
-            includes += ['{}/hiprand/include'.format(ROCM_HOME), '{}/rocrand/include'.format(ROCM_HOME)]
         return includes
diff --git a/op_builder/transformer.py b/op_builder/transformer.py
index 893145d44d94..8db30fdc6791 100644
--- a/op_builder/transformer.py
+++ b/op_builder/transformer.py
@@ -33,7 +33,4 @@ def sources(self):
 
     def include_paths(self):
         includes = ['csrc/includes']
-        if self.is_rocm_pytorch():
-            from torch.utils.cpp_extension import ROCM_HOME
-            includes += ['{}/hiprand/include'.format(ROCM_HOME), '{}/rocrand/include'.format(ROCM_HOME)]
         return includes