[ROCm] support benchmark flag for MIOpen (#77438)

Fixes #68172.  Generally, this corrects multiple flaky convolution unit test behavior seen on ROCm.

The MIOpen integration has been forcing benchmark=True when calling `torch._C._set_cudnn_benchmark(False)`, typically called by `torch.backends.cudnn.set_flags(enabled=True, benchmark=False)`.  We now add support for MIOpen immediate mode to avoid benchmarking during MIOpen solution selection.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77438
Approved by: https://github.com/ngimel, https://github.com/malfet
This commit is contained in:
Jeff Daily
2022-05-23 17:10:24 +00:00
committed by PyTorch MergeBot
parent cac16e2ee2
commit 9aed30d3ad
3 changed files with 292 additions and 66 deletions

View File

@ -508,12 +508,6 @@ PyObject *THPModule_setBenchmarkCuDNN(PyObject *_unused, PyObject *arg)
{
THPUtils_assert(PyBool_Check(arg), "set_benchmark_cudnn expects a bool, "
"but got %s", THPUtils_typename(arg));
#if defined(USE_ROCM)
if (arg == Py_False) {
TORCH_WARN_ONCE("Disabling benchmark mode for MIOpen is NOT supported. Overriding value to True");
arg = Py_True;
}
#endif
at::globalContext().setBenchmarkCuDNN(arg == Py_True);
Py_RETURN_NONE;
}