[TF:TRT] Create execution context with device memory if shape output is present in TRT 7

[tfeher]

Workaround for TRT 7 to handle shape outputs.

Previously, we use createExecutionContextWithoutDeviceMemory to create IExecutionContext, so that TF-TRT can manage the needed device memory to support the execution using the IExecutionContext. This triggers a bug in TRT 7. To workaround the bug, we switch to use createExecutionContext for creating IExecutionContext when the TRTEngine has int32 output.

Add two test cases.

Tracker #45481

Tagging @bixia1 for review.

[bixia1]

Looks good to me. I modified the PR description with more details. Please check.

[tfeher]

Thanks @bixia1 for the review, I have updated the comment. Also checked the modified PR description, it looks fine.

[christopherbate]

Based on the code then it's not the fact that there is a shape tensor, but that it is int32 data type. That's not the same thing, right? Because shape tensor has additional qualifications of being 0-d or 1-d.

Does this same behavior occur when index output of "topK" operation is set as output?

Is there no op that we can insert at the output to force TRT into correct behavior?

[tfeher]

Thanks @christopherbate for your comment! You are right, we could and should be more specific when checking the output tensor. I have added checks to ensure that it is 0-d or 1-d.

Does this same behavior occur when index output of "topK" operation is set as output?

The error happens at an assertion of a shape tensor size: "Assertion failed: size == shapeTensorVolume(out.extent)"
I expect that the problem would only appear with int32 outputs that are shape tensors.

Is there no op that we can insert at the output to force TRT into correct behavior?

I am not aware of any.

[bixia1]

I am not sure whether we can/should add check for 0-d or 1-d to filter the cases. It depends on the situation of the bug. For example, if we have this in the network:
shapeOp(4-d) -> reshape(2x2)
Do we need to allocate context with GPU memory to make this work?

[tfeher]

@bixia1, you are right. I have checked it, and the error happens if I insert the reshape. I will revert the change, which filters for the 0 and 1d cases.

I will also add two unit tests that would trigger the bug. To add the tests, I will need to rebase this PR to #52181.

[tfeher]

Rebased to avoid conflicts. Ready to proceed @bixia1

[bixia1]

Looks good to me, thanks for working on this!
Just two minor comments and it should be good for merging after this.

[tfeher]

Thanks @bixia for the review, I have addressed the issues.