Questions about Value_Tensor_map and Evaluated_Value_map? (Not an issue, just try to understand them...)

[maxyanghu]

I have just gone through TRTorch's 2020 GTC talk/slides/documentation focusing mainly on the graph conversion implementation part. There are some confusions of concepts and questions:

1. What's the relationship between torch::jit::Values and torch::jit::IValue, Are they the same thing? I noticed they are used interchangeably in some situations and are referring to different classes in others.
2. Why do we need to record Value -> ITensor map and Value->IValue map? What's the main use of these two maps?

Could someone help me? Thanks in advance!

[narendasan]

(Moving this to discussions so it is easier to find later)

To answer 1.

torch::jit::Values are abstract representations of data flow in the PyTorch graph.

torch::jit::IValues contain realized data which corresponds to the values in the graph. In the JIT interpreter these IValues are put into a stack machine to feed node executors.

So in this case the graph says Values %0 and %1 are just some int, then when executed those Values are represented in the stack by IValues containing 2 and 4, the consumer of %2 will take 6 later in the execution.

To answer 2:

Nodes dont have attributes or anything like ONNX, they are mostly purely functional. All configurations are inputs to the node. So that for say aten::max_pool2d, you would get inputs like this

aten::max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], int[2] dilation=[1, 1], bool ceil_mode=False) -> (Tensor)

So in the compiler we need the realized data (IValues) since TensorRT engines are built ahead of time. The role of evaluators is to calculate these IValues. So kernel size, stride, padding, dilation and ceil mode will all be calculated during the course of compilation. We store the IValues in evaluated_value_map keyed to the torch::jit::Value it would realize in the JIT interpreter because the JIT graph is SSA (Static single assignment) and this new IValue might be used in multiple places.

In TensorRT, the only dataflow is the actual Tensors being operated on, they are represented by nvinfer1::ITensors in the TRT graph.

In the case of aten::max_pool2d the input tensor self and the output tensor need to be represented in the TRT graph. Converters are responsible for creating TRT layers, so they take the input ITensor created by a different layer and all of the IValues calculated by the evaluators to create the new layer. In that process they associate the output JIT Value with the corresponding output TRT ITensor in the value_tensor_map.

So as we traverse the JIT graph during compilation we know that all preceding JIT Values will have either an IValue in the evaluated_value_map or an ITensor in the value_tensor_map. So when we hit max_pool2d or any other op all we need to do is use the input Values to that node to look up in these maps all of the data we need to create the next layer.

If you look at the documentation for how to write a converter you can see how that is used.
https://nvidia.github.io/TRTorch/contributors/writing_converters.html