Function disable_cudnn_autotune

tensorflow/python/framework/test_util.py:1568–1610

Disable autotuning during the call to this function. Some tests want to base assertions on a graph being isomorphic with a copy. To ensure this, this decorator disables autotuning. Args: func: Function to run with CuDNN autotuning turned off. Returns: Decorated function.

(func)

def disable_cudnn_autotune(func):
  """Disable autotuning during the call to this function.

  Some tests want to base assertions on a graph being isomorphic with a copy.
  To ensure this, this decorator disables autotuning.

  Args:
    func: Function to run with CuDNN autotuning turned off.

  Returns:
    Decorated function.
  """

  def decorator(f):

    def decorated(self, *args, **kwargs):
      original_tf_cudnn_use_autotune = os.environ.get("TF_CUDNN_USE_AUTOTUNE")
      os.environ["TF_CUDNN_USE_AUTOTUNE"] = "false"
      original_xla_flags = os.environ.get("XLA_FLAGS")
      new_xla_flags = "--xla_gpu_autotune_level=0"
      if original_xla_flags:
        new_xla_flags = original_xla_flags + " " + new_xla_flags
      os.environ["XLA_FLAGS"] = new_xla_flags

      result = f(self, *args, **kwargs)

      if (original_tf_cudnn_use_autotune is None):
        del os.environ["TF_CUDNN_USE_AUTOTUNE"]
      else:
        os.environ["TF_CUDNN_USE_AUTOTUNE"] = original_tf_cudnn_use_autotune
      if (original_xla_flags is None):
        del os.environ["XLA_FLAGS"]
      else:
        os.environ["XLA_FLAGS"] = original_xla_flags

      return result

    return decorated

  if func is not None:
    return decorator(func)

  return decorator


# The description is just for documentation purposes.