From 7b480869aebce77beb71fd45dc9d14eb1f98ba04 Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp@nvidia.com>
Date: Tue, 20 Aug 2024 11:26:48 -0700
Subject: [PATCH] Fix segmentaiton fault

---
 mlir-tensorrt/executor/lib/CAPI/Runtime/Runtime.cpp   | 11 +++++++++--
 .../python/bindings/Runtime/RuntimePyBind.cpp         |  3 ++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/mlir-tensorrt/executor/lib/CAPI/Runtime/Runtime.cpp b/mlir-tensorrt/executor/lib/CAPI/Runtime/Runtime.cpp
index 49e750ab5..793151dcc 100644
--- 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
index 5a4d10096..9dc26e412 100644
--- 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 =