Fixes call to __call__ method for CallableFunc subclasses (#98)

Co-authored-by: Lisa Ong <[email protected]>

Author: kernhanda

hatlib/cuda_loader.py

@@ -205,11 +205,11 @@ def cleanup_runtime(self, benchmark: bool, working_dir: str):
         cuda.cuCtxDestroy(self.context)
 
     def init_main(self, benchmark: bool, warmup_iters=0, device_id: int = 0, args=[]):
-        self.func_info.verify(args[0])
+        self.func_info.verify(args[0] if benchmark else args)
         self.device_mem = allocate_cuda_mem(self.func_info.arguments)
 
         if not benchmark:
-            transfer_mem_host_to_cuda(device_args=self.device_mem, host_args=args[0], arg_infos=self.func_info.arguments)
+            transfer_mem_host_to_cuda(device_args=self.device_mem, host_args=args, arg_infos=self.func_info.arguments)
 
         self.ptrs = device_args_to_ptr_list(self.device_mem)
 

hatlib/host_loader.py

@@ -112,24 +112,28 @@ def cleanup_runtime(self, benchmark: bool, working_dir: str):
 
 
     def init_main(self, benchmark: bool, warmup_iters=0, device_id: int = 0, args=[]):
-        self.func_info.verify(args[0])
+        self.func_info.verify(args[0] if benchmark else args)
 
         for _ in range(warmup_iters):
-            for arg in args:
-                self.timer_func(*arg, self.timing_arg_val)
+            if benchmark:
+                for arg in args:
+                    self.timer_func(*arg, self.timing_arg_val)
+            else:
+                self.timer_func(*args, self.timing_arg_val)
 
     def main(self, benchmark: bool, iters=1, batch_size=1, min_time_in_sec=0, args=[]) -> Tuple[float, float]:
         batch_timings_ms: List[float] = []
         i = 0
-        i_max = len(args)
+        i_max = len(args) if benchmark else 1
         iterations = 1
         min_time_in_ms = min_time_in_sec * 1000
 
         while sum(batch_timings_ms) < min_time_in_ms and len(batch_timings_ms) < batch_size:
             self.timing_arg_val.value = np.zeros((1,))
 
             for _ in range(iters):
-                self.timer_func(*args[i], self.timing_arg_val)
+                func_args = args[i] if benchmark else args
+                self.timer_func(*func_args, self.timing_arg_val)
                 i = iterations % i_max
                 iterations += 1
 

hatlib/rocm_loader.py

@@ -123,11 +123,11 @@ def cleanup_runtime(self, benchmark: bool, working_dir: str):
         pass
 
     def init_main(self, benchmark: bool, warmup_iters=0, device_id: int = 0, args=[]):
-        self.func_info.verify(args[0])
+        self.func_info.verify(args[0] if benchmark else args)
         self.device_mem = allocate_rocm_mem(benchmark, self.func_info.arguments, device_id)
 
         if not benchmark:
-            transfer_mem_host_to_rocm(device_args=self.device_mem, host_args=args[0], arg_infos=self.func_info.arguments)
+            transfer_mem_host_to_rocm(device_args=self.device_mem, host_args=args, arg_infos=self.func_info.arguments)
 
         class DataStruct(ctypes.Structure):
             _fields_ = [(f"arg{i}", ctypes.c_void_p) for i in range(len(self.func_info.arguments))]