add NVFuser (#260)

pytorch · Jun 3, 2021 · a84c5db · a84c5db
1 parent d86e206
commit a84c5db
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/conftest.py b/conftest.py
@@ -31,6 +31,20 @@ def set_fuser(fuser):
  torch._C._jit_override_can_fuse_on_cpu(False)
  torch._C._jit_override_can_fuse_on_gpu(True)
  torch._C._jit_set_texpr_fuser_enabled(True)
+ elif fuser == "nvfuser":
+ os.environ['PYTORCH_CUDA_FUSER_DISABLE_FALLBACK'] = '1'
+ os.environ['PYTORCH_CUDA_FUSER_DISABLE_FMA'] = '1'
+ os.environ['PYTORCH_CUDA_FUSER_JIT_OPT_LEVEL'] = '0'
+ torch._C._jit_set_texpr_fuser_enabled(False)
+ torch._C._jit_set_profiling_executor(True)
+ torch._C._jit_set_profiling_mode(True)
+ torch._C._jit_can_fuse_on_cpu()
+ torch._C._jit_can_fuse_on_gpu()
+ torch._C._jit_override_can_fuse_on_cpu(False)
+ torch._C._jit_override_can_fuse_on_gpu(False)
+ torch._C._jit_set_nvfuser_guard_mode(True)
+ torch._C._jit_set_nvfuser_enabled(True)
+
 
 def pytest_sessionstart(session):
  try: