[mlir-tensorrt] Add support for non-DPS calling convention

[jhalakpatel]

This commit introduces support for a non-Destination-Passing Style (non-DPS) calling convention in mlir-tensorrt, while maintaining the existing DPS-style interface.

The changes aim to allow users to compile and execute a mlir-tensorrt executable without allocating output memrefs in advance. Removing the output memref allocation restriction alleviates users from computing output shapes in advance and allocating output memrefs. This is critical since it is part of the performance-critical execution loop.

Deferred output allocation has an added advantage as we no longer need to allocate shape upper bound output buffers and copy the exact output buffer from TensorRT results into mlir-tensorrt output buffers.

This approach also allows us to support data-dependent shapes since the outputs are not required to be allocated before execution.

The non-DPS style calling convention implementation leverages the nvinfer1::IOutputAllocator interface for deferred output allocation. Interface functions such as reallocateOutputAsync and notifyShapes store the allocated output buffer address and its shape.

User call sites with existing DPS-style calling convention:

        out_args = ...  # Requires output shape calculation and output buffer allocation before execution.
        session.execute_function(
            "main", in_args=in_args, out_args=out_args, stream=self.stream._active_cuda_stream
        )

Now, we can do the following with non-DPS style convention:

        outputs = session.execute_function(
            "main", in_args=in_args, stream=self.stream._active_cuda_stream
        )

Key changes include:

1. Update the Plan dialect to implement non-DPS calling convention:
   • Updated PlanAllocTensorsPass and CreateClosedRegionsPass to handle non-DPS convention
   • Update Plan transformation to convert DPS and non-DPS group ops to CallAllocOp and CallOp respectively.
2. Enhanced TensorRT runtime support:
   • Modified ConvertTensorRTToTensorRTRuntime to support both calling conventions
   • Added EnqueueAllocOp for non-DPS style execution
   • Added conversion from EnqueueAllocOp to executor CallOp.
   • Implemented OutputAllocator and CustomTensorRTOuputAllocator classes with proper lifetime management using OutputAllocatorTracker.
3. Updated API and runtime interfaces:
   • Modified executeFunctionWithLuaBackend to support both DPS and non-DPS styles
   • Updated Python bindings to allow non-DPS styles.

[christopherbate]

We should also be verifying the layout (stride + offset) information on memref results, e.g.

trtrt.alloc_enqueue .... -> memref<?x?x?xf32>

implies the identity layout (canonical row major strides).

trtrt.alloc_enqueue .... -> memref<?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>

indicates that the strides are unknown.

Not being able to know anything about the strides is very worst-case since it disables many possible optimizations.

We need this information from TensorRT -- what layouts of results are possible/allowed? Can we enforce that canonical strides will always be returned from TensorRT using the output allocator? If we can, then the verifier should enforce that canonical layouts are used.

Currently for trtrt.enqueue we are effectively enforcing canonical strides for input and output buffers.

[jhalakpatel]

As I mentioned in the other thread, MLIR-TRT does not use the nvinfer1::setAllowedFormat API, which allows formats other than nvinfer1::TensorFormat::kLINEAR. So, for all practical purposes, stride here is canonical.
I will add a check in the verifier for canonical strides. Let me know if you meant this as I might have misunderstood you.

Also, How can I generate this assembly format: memref<?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>?

[jhalakpatel]

This PR will eventually be closed as it has been broken down to smaller PRs:

Non-DPS style functional changes are open for review. I have broken them down to more manageable PR's for review:

1. Output allocator in TensorRT Module: [Lua/TensorRTModule] Implement OutputAllocator for TensorRT execution #282
2. Some improvements to variable names: [NFC] Improve various variable names #283
3. Adding an option to dump textual pipeline: [Compiler] Split compiler pipeline into pass pipeline, compilation, and translation steps #284
4. Update Plan dialect to support non-DPS calling convention: [Dialect/Plan] Update Plan dialect to use non-dps calling convention #285
5. Update TensorRT runtime dialect to use non-dps calling convention: [Dialect/TensorRTRuntime] Update tensorrt runtime dialect to use non-dps calling convention #286
6. Update TensorRT dialect: [Dialect/TensorRT] Update tensorrrt dialect to use non-dps calling convention #287
7. Update Lua Runtime: [Runtime/Lua] Support multiple results in executeFunctionWithLuaBackend #288