Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement custom memory management for internal and output allocations #94

Draft
wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

Implement custom memory management for internal and output allocations #94

wants to merge 6 commits into from

Conversation

jhalakpatel
Copy link
Collaborator

No description provided.

jhalakpatel
jhalakpatel commented Aug 13, 2024
View reviewed changes
mlir-tensorrt/tensorrt/lib/Target/TranslateToTensorRT.cpp Outdated
std::unique_ptr<nvinfer1::IGpuAllocator> gpuAllocator(
new TensorRTLinearAllocator(kCURR_ALLOC_SIZE));
if (opts.useCustomGpuAllocator) {
builder->setGpuAllocator(gpuAllocator.get());
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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 jhalakpatel force-pushed the jhalakp-igpu-allocator branch 5 times, most recently from 842ad29 to dec20fe Compare August 15, 2024 22:54
jhalakpatel
jhalakpatel commented Aug 15, 2024
View reviewed changes
mlir-tensorrt/executor/include/mlir-executor/Runtime/API/API.h Outdated
private:
RuntimeSessionOptions options;
ExecutableView executable;

std::unique_ptr<PinnedMemoryAllocator> pinnedMemoryAllocator;
std::unique_ptr<AllocTracker> allocTracker;
std::unique_ptr<ResourceTracker> resourceTracker;

GpuAllocator* gpuAllocator;
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is unclear how to manage the lifetime. Should I use unique_ptr instead of the raw pointer?

jhalakpatel
jhalakpatel commented Aug 15, 2024
View reviewed changes
mlir-tensorrt/executor/include/mlir-executor/Runtime/Backend/Lua/LuaRuntime.h Outdated
@@ -36,7 +36,7 @@ namespace mlirtrt::runtime {
/// `main` function. It is assumed that `main` takes no arguments and returns an
/// 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);
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it OK to add pass allocator here?

jhalakpatel
jhalakpatel commented Aug 15, 2024
View reviewed changes
mlir-tensorrt/executor/lib/Runtime/Backend/Lua/Modules/TensorRT/TensorRTModule.cpp Outdated
//! 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 {
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I might add this in a separate PR.

@jhalakpatel jhalakpatel force-pushed the jhalakp-igpu-allocator branch 3 times, most recently from 1ddb4b0 to 7b9a0f7 Compare August 16, 2024 22:10
jhalakpatel
jhalakpatel commented Aug 19, 2024
View reviewed changes
mlir-tensorrt/executor/include/mlir-executor/Support/Allocators.h Outdated
class GpuAllocator {
public:
GpuAllocator() = default;
virtual ~GpuAllocator() = default;

virtual StatusOr<void *> reallocate(void *baseAddr, uint64_t alignment,
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Temporarily remove it to simplify testing.

@jhalakpatel jhalakpatel force-pushed the jhalakp-igpu-allocator branch 3 times, most recently from 7b48086 to 0e094ec Compare August 20, 2024 20:28
@jhalakpatel
Copy link
Collaborator Author

ninja -C build/mlir-tensorrt check-mlir-executor

ninja: Entering directory `build/mlir-tensorrt'
[0/2] Re-checking globbed directories...
[0/2] Running the mlir-executor regression tests

Testing Time: 1.60s

Total Discovered Tests: 26
  Unsupported:  1 (3.85%)
  Passed     : 25 (96.15%)

ninja -C build/mlir-tensorrt check-mlir-tensorrt-dialect

ninja: Entering directory `build/mlir-tensorrt'
[0/2] Re-checking globbed directories...
[0/2] Running the mlir-tensorrt-dialect regression tests

Testing Time: 0.10s

Total Discovered Tests: 16
  Unsupported:  1 (6.25%)
  Passed     : 15 (93.75%)

ninja -C build/mlir-tensorrt check-mlir-tensorrt

ninja: Entering directory `build/mlir-tensorrt'
[0/2] Re-checking globbed directories...
[2/4] Running the mlir-tensorrt regression tests
Parallelism Groups: {'non-collective': 5, 'collective': 1, 'models': 1}
llvm-lit: /workspaces/TensorRT-Incubator/mlir-tensorrt/llvm-project/llvm/utils/lit/lit/discovery.py:250: warning: test suite 'MLIR-TensorRT-Unit' contained no tests

Testing Time: 10.97s

Total Discovered Tests: 193
  Unsupported: 99 (51.30%)
  Passed     : 94 (48.70%)

1 warning(s) in tests

@jhalakpatel jhalakpatel force-pushed the jhalakp-igpu-allocator branch 5 times, most recently from 06102b6 to 94cbb5a Compare August 22, 2024 20:46
@jhalakpatel jhalakpatel force-pushed the jhalakp-igpu-allocator branch from 94cbb5a to fc7ca69 Compare August 23, 2024 23:33
@jhalakpatel jhalakpatel changed the title Add IGpuAllocator to MLIR-TensorRT Implement custom memory management for internal and output allocations Aug 26, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@christopherbate christopherbate Awaiting requested review from christopherbate christopherbate will be requested when the pull request is marked ready for review christopherbate is a code owner

@shelkesagar29 shelkesagar29 Awaiting requested review from shelkesagar29 shelkesagar29 will be requested when the pull request is marked ready for review shelkesagar29 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@jhalakpatel