Implement custom memory management for internal and output allocations #94
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jhalakpatel
commented
Aug 13, 2024
| std::unique_ptr<nvinfer1::IGpuAllocator> gpuAllocator( | ||
| new TensorRTLinearAllocator(kCURR_ALLOC_SIZE)); | ||
| if (opts.useCustomGpuAllocator) { | ||
| builder->setGpuAllocator(gpuAllocator.get()); |
Collaborator
Author
There was a problem hiding this comment.
@christopherbate Here are the builder-side changes to allow the GPU allocator to be set up. We probably do not need this and will only be required for runtime allocations.
jhalakpatel
force-pushed
the
jhalakp-igpu-allocator
branch
5 times, most recently
from
842ad29 to
dec20fe
Compare
jhalakpatel
commented
Aug 15, 2024
| std::unique_ptr<AllocTracker> allocTracker; | ||
| std::unique_ptr<ResourceTracker> resourceTracker; | ||
|
|
||
| GpuAllocator* gpuAllocator; |
Collaborator
Author
There was a problem hiding this comment.
It is unclear how to manage the lifetime. Should I use unique_ptr instead of the raw pointer?
jhalakpatel
commented
Aug 15, 2024
| /// integer result (which is returned if the execution is successful). | ||
| /// TODO: this should take a handle to a function for streaming output/errors. | ||
| StatusOr<int64_t> runExecutorLuaScript(std::string_view luaScript); | ||
| StatusOr<int64_t> runExecutorLuaScript(std::string_view luaScript, GpuAllocator* allocator); |
Collaborator
Author
There was a problem hiding this comment.
Is it OK to add pass allocator here?
jhalakpatel
commented
Aug 15, 2024
| //! Class to allocate memory for outputs with data-dependent shapes. The sizes | ||
| //! of those are unknown so pre-allocation is not possible. | ||
| //! | ||
| class OutputAllocator : public nvinfer1::IOutputAllocator { |
Collaborator
Author
There was a problem hiding this comment.
I might add this in a separate PR.
jhalakpatel
force-pushed
the
jhalakp-igpu-allocator
branch
3 times, most recently
from
1ddb4b0 to
7b9a0f7
Compare
jhalakpatel
commented
Aug 19, 2024
| GpuAllocator() = default; | ||
| virtual ~GpuAllocator() = default; | ||
|
|
||
| virtual StatusOr<void *> reallocate(void *baseAddr, uint64_t alignment, |
Collaborator
Author
There was a problem hiding this comment.
Temporarily remove it to simplify testing.
jhalakpatel
force-pushed
the
jhalakp-igpu-allocator
branch
3 times, most recently
from
7b48086 to
0e094ec
Compare
Collaborator
Author
|
ninja -C build/mlir-tensorrt check-mlir-executor ninja -C build/mlir-tensorrt check-mlir-tensorrt-dialect ninja -C build/mlir-tensorrt check-mlir-tensorrt |
jhalakpatel
force-pushed
the
jhalakp-igpu-allocator
branch
5 times, most recently
from
06102b6 to
94cbb5a
Compare
jhalakpatel
force-pushed
the
jhalakp-igpu-allocator
branch
from
94cbb5a to
fc7ca69
Compare
jhalakpatel
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.