[ROCm] disable NVFuser (#92182)

In preparation for #89621.

Partial reverts of #82498 and #86369.

Pull Request resolved: #92182
Approved by: https://github.com/davidberard98

.github/workflows/_rocm-test.yml

@@ -118,7 +118,6 @@ jobs:
           SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
           DOCKER_IMAGE: ${{ inputs.docker-image }}
           XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
-          PYTORCH_JIT_ENABLE_NVFUSER: 1
           PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }}
           PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
         timeout-minutes: 270

test/test_jit_cuda_fuser.py

@@ -35,7 +35,7 @@
 
 from typing import List
 
-RUN_NVFUSER = RUN_CUDA
+RUN_NVFUSER = RUN_CUDA and not TEST_WITH_ROCM
 CUDA_MAJOR, CUDA_MINOR = 0, 0
 
 if RUN_NVFUSER and torch.version.cuda is not None:
@@ -5197,8 +5197,18 @@ def test_register_fuser_cpu(self):
             torch._C._jit_set_nvfuser_enabled(True)
             torch._C._jit_set_nvfuser_enabled(False)
 
+    @unittest.skipIf(not RUN_CUDA, "requires CUDA")
+    @unittest.skipIf(not TEST_WITH_ROCM, "ROCM test only")
+    def test_register_fuser_rocm(self):
+        with self.assertRaises(RuntimeError):
+            torch._C._jit_set_nvfuser_enabled(True)
+            torch._C._jit_set_nvfuser_enabled(False)
+
     def test_can_be_enabled_nvfuser(self):
-        expected = RUN_CUDA
+        if TEST_WITH_ROCM:
+            expected = False
+        else:
+            expected = RUN_CUDA
 
         self.assertEqual(expected, torch._C._jit_nvfuser_can_be_enabled())
 

torch/csrc/jit/codegen/cuda/interface.cpp

@@ -43,7 +43,7 @@ class NVFuserEnabler {
 
  public:
   static bool nvfuserCanBeEnabled() {
-#if defined(FBCODE_CAFFE2)
+#if defined(USE_ROCM) || defined(FBCODE_CAFFE2)
     return false;
 #endif
     return at::globalContext().hasCUDA() &&