Fix segmentaiton fault

NVIDIA · Aug 20, 2024 · 7b48086 · 7b48086
1 parent a7a1ae0
commit 7b48086
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/mlir-tensorrt/executor/lib/CAPI/Runtime/Runtime.cpp b/mlir-tensorrt/executor/lib/CAPI/Runtime/Runtime.cpp
@@ -668,6 +668,11 @@ class GpuAllocatorWrapper : public GpuAllocator {
 
   // Static method to create a GpuAllocator from MTRT_GpuAllocator
   static std::unique_ptr<GpuAllocator> create(MTRT_GpuAllocator gpuAllocator) {
+    if (!gpuAllocator.ptr || !gpuAllocator.allocate ||
+        !gpuAllocator.deallocate) {
+      llvm::errs() << "Invalid MTRT_GpuAllocator passed to create()";
+      return nullptr;
+    }
     return std::make_unique<GpuAllocatorWrapper>(gpuAllocator);
   }
 };
@@ -679,8 +684,10 @@ MTRT_Status mtrtRuntimeSessionCreate(MTRT_RuntimeSessionOptions options,
   RuntimeSessionOptions *cppOptions = unwrap(options);
   Executable *cppExecutable = unwrap(executable);
 
-  std::unique_ptr<GpuAllocator> allocator =
-      gpuAllocator.ptr ? GpuAllocatorWrapper::create(gpuAllocator) : nullptr;
+  std::unique_ptr<GpuAllocator> allocator;
+  if (gpuAllocator.ptr) {
+    allocator.reset(GpuAllocatorWrapper::create(gpuAllocator).release());
+  }
 
   StatusOr<std::unique_ptr<RuntimeSession>> session =
       createRuntimeSessionWithLuaBackend(cppExecutable->getView(),

diff --git a/mlir-tensorrt/python/bindings/Runtime/RuntimePyBind.cpp b/mlir-tensorrt/python/bindings/Runtime/RuntimePyBind.cpp
@@ -983,7 +983,8 @@ PYBIND11_MODULE(_api, m) {
              if (gpu_allocator.is_none()) {
                // Create session without custom allocator
                s = mtrtRuntimeSessionCreate(
-                   options, exe, MTRT_GpuAllocator{nullptr}, &session);
+                   options, exe, MTRT_GpuAllocator{nullptr, nullptr, nullptr},
+                   &session);
              } else {
                try {
                  PyGpuAllocator &allocator =