[TENSORRT] Improvements and fixes for TensorRT

[mbaret]

A number of small fixes and refactors to improve the robustness of the TensorRT integration.

Co-authored-by: Mark Shields [email protected]

[mbaret]

cc @mbs-octoml @mikepapadim

[mbs-octoml]

The background on the bits I fiddled with:

• Finish the transition from operator predicate based to pattern based for the partition_for_tensorrt function. The current state in main is broken.
• Remove choice of operator vs predicate method since they are indistinguishable from the outside.
• Though most operators can be translated looking only at the operator call, nn.batch_norm can only be translated in the form of a sub-graph nn.batch_norm(...).0. So switch the translation to follow the standard composite style.

[mbs-octoml]

LGTM, thank you!

I wonder if we should warn for the conv3d attributes you found trigger numerical inaccuracy? Say in the predicate?

[mbaret]

Unfortunately, the attributes aren't particularly well-defined. For example, in that specific test it was failing with channels between 28 and 45. I think so long as TensorRT is selecting strategies with different numerical behaviour based on tuning we won't be able to give a sane warning here.

[mbs-octoml]

Ack. We do what we can.

[mikepapadim]

LGTM, thanks

[jwfromm]

Thanks for these improvements @mbaret, @mikepapadim, and @mbs-octoml!

[tiandiao123]

nice job!