diff --git a/.github/workflows/_rocm-test.yml b/.github/workflows/_rocm-test.yml index d600ba314007b..93b70aa97a73c 100644 --- a/.github/workflows/_rocm-test.yml +++ b/.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 diff --git a/test/test_jit_cuda_fuser.py b/test/test_jit_cuda_fuser.py index 0a13fdb20a823..4f0ea9dcd3443 100644 --- a/test/test_jit_cuda_fuser.py +++ b/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()) diff --git a/torch/csrc/jit/codegen/cuda/interface.cpp b/torch/csrc/jit/codegen/cuda/interface.cpp index 2a9fa2f5dffce..12126726aa4d7 100644 --- a/torch/csrc/jit/codegen/cuda/interface.cpp +++ b/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() &&