Allow disabling nvfuser without CUDA (#71358)

Summary:
On a CPU-only build of pytorch `torch._C._jit_set_nvfuser_enabled(False)` would throw an error (even though it is a no-op operation), with this fix:
```
>>> torch._C._jit_set_nvfuser_enabled(False)
False
>>> torch._C._jit_set_nvfuser_enabled(True)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Running CUDA fuser is only supported on CUDA builds.
>>>
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/71358

Reviewed By: eellison

Differential Revision: D33601135

Pulled By: jansel

fbshipit-source-id: c764df2fa197ce7b4f71e5df0a91cd988766e99c
(cherry picked from commit a801df9321)
This commit is contained in:
Jason Ansel 2022-01-19 11:54:24 -08:00 committed by PyTorch MergeBot
parent 214f4bf2ff
commit ac26f8237c

View File

@ -12,11 +12,11 @@ namespace jit {
struct C10_EXPORT RegisterCudaFuseGraph
: public PassManager<RegisterCudaFuseGraph> {
static bool registerPass(bool enabled) {
bool old_flag = PassManager::isRegistered();
if (enabled) {
TORCH_CHECK(
at::globalContext().hasCUDA() && !at::globalContext().hasHIP(),
"Running CUDA fuser is only supported on CUDA builds.");
bool old_flag = PassManager::isRegistered();
if (enabled) {
PassManager::registerPass(fuser::cuda::fuseGraph);
} else {
PassManager::clearPass();