Update tests to use custom allocator

jhalakpatel · jhalakpatel · commit a7a1ae06d51d · 2024-08-20T11:13:47.000-07:00
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/ClusteringDynamicShape/end-to-end-binary.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/ClusteringDynamicShape/end-to-end-binary.mlir
@@ -3,7 +3,7 @@
 // RUN: stablehlo-clustering-pipeline, \
 // RUN: post-clustering-pipeline, \
 // RUN: executor-lowering-pipeline)" \
-// RUN: | mlir-tensorrt-translate -mlir-to-runtime-executable -allow-unregistered-dialect |  mlir-tensorrt-runner -input-type=rtexe
+// RUN: | mlir-tensorrt-translate -mlir-to-runtime-executable -allow-unregistered-dialect |  mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator
 
 #profile = #tensorrt.shape_profile<min = [2], opt = [4], max = [6]>
 #profile1 = #tensorrt.shape_profile<min = [1], opt = [3], max = [5]>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/ClusteringDynamicShape/end-to-end-unary.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/ClusteringDynamicShape/end-to-end-unary.mlir
@@ -3,7 +3,7 @@
 // RUN: stablehlo-clustering-pipeline, \
 // RUN: post-clustering-pipeline, \
 // RUN: executor-lowering-pipeline)" \
-// RUN: | mlir-tensorrt-translate -mlir-to-runtime-executable -allow-unregistered-dialect |  mlir-tensorrt-runner -input-type=rtexe
+// RUN: | mlir-tensorrt-translate -mlir-to-runtime-executable -allow-unregistered-dialect |  mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator
 
 #profile0 = #tensorrt.shape_profile<min = [1], opt = [5], max = [10]>
 #profile1 = #tensorrt.shape_profile<min = [2], opt = [5], max = [6]>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-bf16.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-bf16.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xbf16, #plan.memory_space<host_pinned>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-dynamic.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-dynamic.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 func.func @run_with_shape_2d(%arg0: memref<?xindex>, %arg1: memref<2xindex>) {
   %c0 = arith.constant 0 : index
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f16.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f16.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xf16, #plan.memory_space<host_pinned>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f32.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f32.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate  -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xf32, #plan.memory_space<host>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f8E4M3FN.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-f8E4M3FN.mlir
@@ -1,6 +1,6 @@
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xf8E4M3FN, #plan.memory_space<host_pinned>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-i1.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-i1.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xi1, #plan.memory_space<host>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-i4.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/buffer-ops-i4.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 !descriptor1D = !executor.table<!executor.ptr<device>, !executor.ptr<device>, index, index, index>
 !hostMemRef = memref<4xi4, #plan.memory_space<host_pinned>>
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/memcpy-strided.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/memcpy-strided.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 func.func @main() -> index {
   %c0 = arith.constant 0 : index
diff --git a/mlir-tensorrt/test/Target/Lua/IntegrationTests/memcpy.mlir b/mlir-tensorrt/test/Target/Lua/IntegrationTests/memcpy.mlir
@@ -1,7 +1,7 @@
 // REQUIRES: host-has-at-least-1-gpus
 // RUN: mlir-tensorrt-opt %s -convert-memref-to-cuda -convert-plan-to-executor -convert-cuda-to-executor -executor-lowering-pipeline \
 // RUN:   | mlir-tensorrt-translate -mlir-to-runtime-executable \
-// RUN:   | mlir-tensorrt-runner -input-type=rtexe | FileCheck %s
+// RUN:   | mlir-tensorrt-runner -input-type=rtexe --use-custom-allocator | FileCheck %s
 
 func.func @main() -> i32 {
   %c0_i32 = arith.constant 0 : i32
diff --git a/mlir-tensorrt/test/python/IntegrationTests/TRT10/test_stablehlo_add.py b/mlir-tensorrt/test/python/IntegrationTests/TRT10/test_stablehlo_add.py
@@ -33,7 +33,7 @@ def test_stablehlo_add(
             exe = compiler.compiler_stablehlo_to_executable(client, m.operation, opts)
 
         session_options = runtime.RuntimeSessionOptions(num_devices=1, device_id=0)
-        session = runtime.RuntimeSession(session_options, exe)
+        session = runtime.RuntimeSession(session_options, exe, None)
 
         session.execute_function(
             "main", in_args=test.in_args, out_args=test.out_args, stream=stream
diff --git a/mlir-tensorrt/test/python/IntegrationTests/test_call_validation.py b/mlir-tensorrt/test/python/IntegrationTests/test_call_validation.py
@@ -62,7 +62,7 @@ def create_scalar(self, value):
         return self.client.create_scalar(value, runtime.ScalarTypeCode.i64)
 
     def execute(self, arg: runtime.RuntimeValue):
-        session = runtime.RuntimeSession(self.session_options, self.exe)
+        session = runtime.RuntimeSession(self.session_options, self.exe, None)
         try:
             session.execute_function(
                 "main", in_args=[arg], out_args=[arg], stream=self.stream
diff --git a/mlir-tensorrt/test/python/IntegrationTests/test_stablehlo_add.py b/mlir-tensorrt/test/python/IntegrationTests/test_stablehlo_add.py
@@ -15,29 +15,26 @@
 """
 
 
-class AppGpuAllocator(runtime.GpuAllocator):
+class CupyGPUAllocator(runtime.GpuAllocator):
     def __init__(self):
-        # Initialize the base class
         super().__init__()
+        self.allocations = {}  # Keep track of allocations
 
-    def allocate(self, size, alignment, flags=None, stream=None):
-        # Implement memory allocation using CuPy
-        # Allocate memory on the GPU
-        import pdb
+    def allocate(self, size):
+        # Allocate memory on the GPU using CuPy
+        mem = cp.cuda.alloc(size)
+        ptr = int(mem.ptr)  # Convert to integer
+        # Store the CuPy memory object
+        self.allocations[ptr] = mem
+        return ptr
 
-        pdb.set_trace()
-        memory = cp.empty(size, dtype=cp.uint8)
-        print(
-            f"Allocated memory: {memory.data.ptr}"
-        )  # Print pointer to allocated GPU memory
-        return memory  # Return the CuPy array
-
-    def deallocate(self, memory, stream=None):
-        # Deallocate memory in CuPy
-        # CuPy handles memory automatically, so explicit deallocation is not necessary
-        print(f"Deallocating memory: {memory}")
-        del memory  # Explicitly delete the CuPy array if needed
-        return True
+    def deallocate(self, ptr):
+        if ptr in self.allocations:
+            # Remove the reference to the CuPy memory object
+            # This will trigger deallocation if there are no other references
+            del self.allocations[ptr]
+            return True
+        return False
 
 
 def stablehlo_add():
@@ -62,7 +59,8 @@ def stablehlo_add():
     if len(devices) == 0:
         return
 
-    allocator = AppGpuAllocator()
+    # Create an instance of the custom allocator
+    allocator = CupyGPUAllocator()
 
     session_options = runtime.RuntimeSessionOptions(num_devices=1, device_id=0)
     session = runtime.RuntimeSession(session_options, exe, gpu_allocator=allocator)
diff --git a/mlir-tensorrt/test/python/IntegrationTests/test_stablehlo_dynamic.py b/mlir-tensorrt/test/python/IntegrationTests/test_stablehlo_dynamic.py
@@ -115,7 +115,7 @@ def test_program(program: str, input_shape: Iterable[int], debug: bool = True):
         return
 
     session_options = runtime.RuntimeSessionOptions(num_devices=1, device_id=0)
-    session = runtime.RuntimeSession(session_options, exe)
+    session = runtime.RuntimeSession(session_options, exe, gpu_allocator=None)
 
     arg0 = client.create_memref(
         np.ones(input_shape, dtype=np.float32).data,
