Fix segmentaiton fault

Author: jhalakpatel

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(),

mlir-tensorrt/python/bindings/Runtime/RuntimePyBind.cpp

@@ -191,6 +191,9 @@ class PyRuntimeValue : public PyMTRTWrapper<PyRuntimeValue, MTRT_RuntimeValue> {
 // routing.
 class PyGpuAllocator {
 public:
+  py::object pySelf;
+  PyGpuAllocator(py::object self) : pySelf(self) {}
+
   virtual ~PyGpuAllocator() = default;
   virtual std::uintptr_t allocate(uint64_t size) = 0;
   virtual bool deallocate(std::uintptr_t ptr) = 0;
@@ -201,6 +204,7 @@ class PyGpuAllocator {
 private:
   // Trampoline function: Routes C-style allocation calls to C++ virtual method.
   static void *pyGpuAllocatorAllocate(void *self, uint64_t size) {
+    py::gil_scoped_acquire acquire;
     auto *allocator = static_cast<PyGpuAllocator *>(self);
     std::uintptr_t ptr = allocator->allocate(size);
     return reinterpret_cast<void *>(ptr);
@@ -209,6 +213,7 @@ class PyGpuAllocator {
   // Trampoline function: Routes C-style deallocation calls to C++ virtual
   // method.
   static bool pyGpuAllocatorDeallocate(void *self, void *memory) {
+    py::gil_scoped_acquire acquire;
     auto *allocator = static_cast<PyGpuAllocator *>(self);
     return allocator->deallocate(reinterpret_cast<std::uintptr_t>(memory));
   }
@@ -969,7 +974,8 @@ PYBIND11_MODULE(_api, m) {
            py::arg("nccl_uuid") = py::str(""));
 
   py::class_<PyGpuAllocator, PyGpuAllocatorTrampoline>(m, "GpuAllocator")
-      .def(py::init<>())
+      .def(py::init<>(
+          [](py::object self) { return new PyGpuAllocatorTrampoline(self); }))
       .def("allocate", &PyGpuAllocator::allocate)
       .def("deallocate", &PyGpuAllocator::deallocate)
       .def("get_capi_object", &PyGpuAllocator::getCApiObject);
@@ -983,7 +989,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 =

mlir-tensorrt/test/python/IntegrationTests/test_stablehlo_add.py

@@ -17,7 +17,7 @@
 
 class CupyGPUAllocator(runtime.GpuAllocator):
     def __init__(self):
-        super().__init__()
+        super().__init__(self)
         self.allocations = {}  # Keep track of allocations
 
     def allocate(self, size):