Change run_iree_module_function to accept HAL device instead of drive…

…r name (#893)

It did not sit right that run_iree_module_function would create a HAL
device just to make iree.runtime.FunctionInvoker happy.

sharktank/sharktank/utils/iree.py

@@ -144,7 +144,7 @@ def run_iree_module_function(
     module: iree.runtime.VmModule,
     vm_context: iree.runtime.VmContext,
     args: List[iree.runtime.DeviceArray],
-    driver: str,
+    device: iree.runtime.HalDevice,
     function_name: str = "main",
     trace_path_prefix: Optional[str] = None,
 ) -> List[iree.runtime.DeviceArray]:
@@ -154,7 +154,7 @@ def run_iree_module_function(
         vm_context=vm_context,
         # TODO: rework iree.runtime.FunctionInvoker interface for multiple devices.
         # This works, but does not look right.
-        device=iree.runtime.get_device(driver, cache=False),
+        device=device,
         vm_function=vm_function,
     )
 

sharktank/tests/models/clip/clip_test.py

@@ -200,7 +200,7 @@ def runTestCompareIreeAgainstTorchEagerWithInputTokens(
                 module=iree_module,
                 vm_context=iree_vm_context,
                 args=iree_args,
-                driver="hip",
+                device=iree_devices[0],
                 function_name=f"forward_bs{batch_size}",
                 trace_path_prefix=f"{target_model_path_prefix}_iree_",
             )

sharktank/tests/models/flux/flux_test.py

@@ -160,7 +160,7 @@ def runCompareIreeAgainstTorchEager(
                 module=iree_module,
                 vm_context=iree_vm_context,
                 args=iree_args,
-                driver="hip",
+                device=iree_devices[0],
                 function_name=f"forward_bs{batch_size}",
             )
         )

sharktank/tests/models/llama/sharded_llama_test.py

@@ -340,7 +340,7 @@ def _(model, *args, **kwargs) -> torch.Tensor:
             function_name="prefill",
             module=iree_module,
             vm_context=vm_context,
-            driver=iree_driver,
+            device=iree_devices[0],
             trace_path_prefix=path_prefix if dump_enabled else None,
         )
         prefill_iree_result = UnreducedTensor(ts=iree_to_torch(*prefill_iree_result))
@@ -367,7 +367,7 @@ def _(model, *args, **kwargs) -> torch.Tensor:
             function_name="decode",
             module=iree_module,
             vm_context=vm_context,
-            driver=iree_driver,
+            device=iree_devices[0],
             trace_path_prefix=path_prefix if dump_enabled else None,
         )
         decode_iree_result = UnreducedTensor(ts=iree_to_torch(*decode_iree_result))

sharktank/tests/models/t5/t5_test.py

@@ -365,7 +365,7 @@ def runTestV1_1CompareIreeAgainstTorchEager(
                 module=iree_module,
                 vm_context=iree_vm_context,
                 args=iree_args,
-                driver="hip",
+                device=iree_devices[0],
                 function_name=f"forward_bs{batch_size}",
                 trace_path_prefix=f"{target_model_path_prefix}_iree_",
             )

sharktank/tests/models/vae/vae_test.py

@@ -165,7 +165,7 @@ def testVaeIreeVsHuggingFace(self):
             module=iree_module,
             vm_context=iree_vm_context,
             args=iree_args,
-            driver="hip",
+            device=iree_devices[0],
             function_name="decode",
         )[0].to_host()
         # TODO: Verify these numerics are good or if tolerances are too loose
@@ -193,7 +193,7 @@ def testVaeIreeVsHuggingFace(self):
             module=iree_module,
             vm_context=iree_vm_context,
             args=iree_args,
-            driver="hip",
+            device=iree_devices[0],
             function_name="decode",
         )[0].to_host()
         # TODO: Upload IR on passing tests
@@ -328,7 +328,7 @@ def testVaeIreeVsHuggingFace(self):
                 module=iree_module,
                 vm_context=iree_vm_context,
                 args=iree_args,
-                driver="hip",
+                device=iree_devices[0],
                 function_name="decode",
             )[0]
         )
@@ -355,7 +355,7 @@ def testVaeIreeVsHuggingFace(self):
                 module=iree_module,
                 vm_context=iree_vm_context,
                 args=iree_args,
-                driver="hip",
+                device=iree_devices[0],
                 function_name="decode",
             )[0]
         )

sharktank/tests/ops/ops_test.py

@@ -384,7 +384,7 @@ def _(model, x):
             module=iree_module,
             vm_context=iree_vm_context,
             args=iree_args,
-            driver="local-task",
+            device=iree_devices[0],
             function_name=f"forward",
         )