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

[hunbssfy]

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

[hunbssfy]

Thanks for giving a very thorough explanation!
About this part:

So in the compiler we need the realized data (IValues) since TensorRT engines are built ahead of time

Why use torch::jit::IValue instead of defining a new class for storing static values? Are you using JIT mechanism to evaluate these static values?

[narendasan]

IValue is a tagged union of all possible data types that PyTorch supports (like enums in rust) so it makes it very easy to move data around in a generic wrapper. We thought it was much easier to leverage a lot of PyTorch's infrastructure instead of re-implementing since we already have the dependency on libtorch. We don't use JIT right now the evaluate static values, we create a library of evaluators which do the same calculations, though it might be interesting to see if we can reuse PyTorch for that.

[hunbssfy]

@narendasan Thanks.
Another question is about condition block compilation: I read the code and saw this line in EvaluateConditionalBlock:
auto condition = ctx->evaluated_value_map[n->input(0)].toBool();
What is input(0) is a tensor that cannot be evaluated at compile time? Is there a handle for this scenario?

[narendasan]

The code here directly evaluates according to TorchScript conditional semantics https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/OVERVIEW.md#if. The n->input(0) shouldnt be anything but a bool according to the semantics.