[JIT] Add an environment variable for controlling verbosity of JIT IR dumps #41770
Description
🚀 Feature
As of #41507, JIT IR dumps started to include full type information and that made IR dumps more verbose. Sometimes it is convenient to hide some of the info, it would be nice to have an environment variable controlling it.
Test case
import torch
def bar():
def test(a):
return a
x = torch.ones(10,10, device='cpu')
print(torch.jit.trace(test, (x)).graph)
Current output:
$ python test.py
graph(%0 : Float(10:10, 10:1, device=cpu)):
return (%0)
Desired functionality:
$ PYTORCH_JIT_TYPE_VERBOSITY=0 python test.py # Dont show any types
graph(%0):
return (%0)
$ PYTORCH_JIT_TYPE_VERBOSITY=1 python test.py # Show sizes and dtype in types
graph(%0 : Float(10, 1)):
return (%0)
$ PYTORCH_JIT_TYPE_VERBOSITY=2 python test.py # Show sizes, strides and dtype in types
graph(%0 : Float(10:10, 1:1)):
return (%0)
$ PYTORCH_JIT_TYPE_VERBOSITY=3 python test.py # Show full type info
graph(%0 : Float(10:10, 1:1, device=cpu)):
return (%0)
Implementation pointers
Function that prints JIT IR types is located in aten/src/ATen/core/type.cpp (look for std::ostream& operator<<(std::ostream & out, const Type & t)).
An example of how env variable could be used for controlling logging behavior could be found in torch/csrc/jit/jit_log.cpp (look for PYTORCH_JIT_LOG_LEVEL).
cc @suo @gmagogsfm