Merge pull request #9443 from google/benvanik-device-flags

Replacing `--driver=` with `--device=` and adding list flags.

build_tools/bazel/iree_check_test.bzl

@@ -67,7 +67,7 @@ def iree_check_test(
     native_test(
         name = name,
         args = [
-            "--driver=%s" % driver,
+            "--device=%s" % driver,
             "$(location :%s)" % bytecode_module_name,
         ] + runner_args,
         data = [":%s" % bytecode_module_name],

build_tools/bazel/iree_trace_runner_test.bzl

@@ -65,7 +65,7 @@ def iree_trace_runner_test(
     native_test(
         name = name,
         args = [
-            "--driver=%s" % driver,
+            "--device=%s" % driver,
             "$(location :%s)" % trace,
         ] + runner_args,
         data = [

build_tools/benchmarks/comparisons/common/benchmark_command.py

@@ -129,7 +129,7 @@ def parse_latency_from_output(self, output: str) -> float:
 
   def generate_benchmark_command(self) -> list[str]:
     command = super().generate_benchmark_command()
-    command.append("--driver=" + self.driver)
+    command.append("--device=" + self.driver)
     command.append("--task_topology_group_count=" + str(self.num_threads))
     command.append("--benchmark_repetitions=" + str(self.num_runs))
     return command

build_tools/cmake/iree_benchmark_suite.cmake

@@ -300,7 +300,7 @@ function(iree_benchmark_suite)
         COMMAND
           "${Python3_EXECUTABLE}" "${IREE_ROOT_DIR}/build_tools/scripts/generate_flagfile.py"
             --module_file="${_MODULE_FILE_FLAG}"
-            --driver=${_RULE_DRIVER}
+            --device=${_RULE_DRIVER}
             --entry_function=${_MODULE_ENTRY_FUNCTION}
             --function_inputs=${_MODULE_FUNCTION_INPUTS}
             "${_ADDITIONAL_ARGS_CL}"

build_tools/cmake/iree_native_test.cmake

@@ -14,7 +14,7 @@ include(CMakeParseArguments)
 #
 # Parameters:
 # NAME: name of target
-# DRIVER: If specified, will pass --driver=DRIVER to the test binary and adds
+# DRIVER: If specified, will pass --device=DRIVER to the test binary and adds
 #     a driver label to the test.
 # TEST_INPUT_FILE_ARG: If specified, the input file will be added to DATA and
 #     its device path appended to ARGS. Note that the device path may be
@@ -24,7 +24,7 @@ include(CMakeParseArguments)
 #     a separate device (e.g. Android), these files will be pushed to the
 #     device. TEST_INPUT_FILE_ARG is automatically added if specified.
 # ARGS: additional arguments passed to the test binary. TEST_INPUT_FILE_ARG and
-#     --driver=DRIVER are automatically added if specified.
+#     --device=DRIVER are automatically added if specified.
 # SRC: binary target to run as the test.
 # LABELS: Additional labels to apply to the test. The package path is added
 #     automatically.
@@ -69,7 +69,7 @@ function(iree_native_test)
 
   # If driver was specified, add the corresponding test arg and label.
   if(DEFINED _RULE_DRIVER)
-    list(APPEND _RULE_ARGS "--driver=${_RULE_DRIVER}")
+    list(APPEND _RULE_ARGS "--device=${_RULE_DRIVER}")
     list(APPEND _RULE_LABELS "driver=${_RULE_DRIVER}")
   endif()
 

build_tools/kokoro/gcp_ubuntu/cmake/linux/riscv64/tests/lit.cfg.py

@@ -20,7 +20,7 @@
     (os.getenv("QEMU_RV64_BIN"), os.getenv("RISCV_TOOLCHAIN_ROOT")))
 
 config.environment["TEST_MODULE_CMD"] = (
-    "%s %s/tools/iree-run-module --driver=local-task" %
+    "%s %s/tools/iree-run-module --device=local-task" %
     (config.environment["TEST_CMD"], os.getenv("BUILD_RISCV_DIR")))
 
 config.test_exec_root = os.getenv("BUILD_RISCV_DIR") + \

build_tools/scripts/generate_flagfile.py

@@ -19,11 +19,11 @@ def parse_arguments():
                       required=True,
                       metavar="<module-file>",
                       help="The name of the module file")
-  parser.add_argument("--driver",
+  parser.add_argument("--device",
                       type=str,
                       required=True,
-                      metavar="<driver>",
-                      help="The name of the IREE driver")
+                      metavar="<device>",
+                      help="The name of the HAL device")
   parser.add_argument("--entry_function",
                       type=str,
                       required=True,
@@ -50,7 +50,7 @@ def parse_arguments():
 
 def main(args):
   lines = [
-      f"--driver={args.driver}", f"--module_file={args.module_file}",
+      f"--device={args.device}", f"--module_file={args.module_file}",
       f"--entry_function={args.entry_function}"
   ]
   lines.extend([

docs/developers/design_docs/cuda_backend.md

@@ -85,7 +85,7 @@ $ ../iree-build/tools/iree-compile \
 
 # Run the module through CUDA HAL backend.
 $ ../iree-build/tools/iree-run-module \
-  --driver=cuda \
+  --device=cuda \
   --module_file=/tmp/mhlo-add.vmfb \
   --entry_function=add \
   --function_input="4xf32=[1 2 3 4]" \

docs/developers/developing_iree/benchmarking.md

@@ -31,7 +31,7 @@ and then benchmark an exported function in that module:
 ```shell
 $ bazel run //tools:iree-benchmark-module -- \
   --module_file=/tmp/module.fb \
-  --driver=local-task \
+  --device=local-task \
   --entry_function=abs \
   --function_input=f32=-2
 ```
@@ -78,7 +78,7 @@ Now we'll actually invoke the binary:
 ```shell
 $ ./bazel-bin/tools/iree-benchmark-module \
   --module_file=/tmp/module.fb \
-  --driver=local-task \
+  --device=local-task \
   --entry_function=abs \
   --function_input=f32=-2
 ```
@@ -122,7 +122,7 @@ in that module:
 ```shell
 $ build/tools/iree-benchmark-module
   --module_file=/tmp/fullyconnected.vmfb
-  --driver=local-task
+  --device=local-task
 ```
 
 If no `entry_function` is specified, `iree-benchmark-module` will register a

docs/developers/developing_iree/developer_overview.md

@@ -140,7 +140,7 @@ above on IREE's VMVX driver:
 ```shell
 $ ../iree-build/tools/iree-run-module \
   --module_file=/tmp/simple_abs_vmvx.vmfb \
-  --driver=local-task \
+  --device=local-task \
   --entry_function=abs \
   --function_input=f32=-2
 ```
@@ -165,7 +165,7 @@ $ ../iree-build/tools/iree-compile \
 ```shell
 $ ../iree-build/tools/iree-check-module \
   /tmp/abs.vmfb \
-  --driver=local-task
+  --device=local-task
 ```
 
 ### iree-run-mlir

docs/developers/developing_iree/e2e_benchmarking.md

@@ -117,7 +117,7 @@ on VMVX run:
 ```shell
 $ tools/iree-benchmark-module \
   --module_file=/tmp/iree/modules/MatrixOpsStaticModule/iree_vmvx/compiled.vmfb \
-  --driver=local-task \
+  --device=local-task \
   --entry_function=matmul_lhs_batch \
   --function_input=256x64x32xf32=2 \
   --function_input=32x16xf32=3
@@ -221,7 +221,7 @@ $ adb push /tmp/iree/modules/MatrixOpsStaticModule/iree_vmvx/* \
 ```shell
 $ adb shell /data/local/tmp/iree-benchmark-module \
   --module_file="/data/local/tmp/MatrixOpsStaticModule/iree_vmvx/compiled.vmfb" \
-  --driver=local-task \
+  --device=local-task \
   --entry_function=matmul_lhs_batch \
   --function_input=256x64x32xf32=2 \
   --function_input=32x16xf32=3

docs/developers/developing_iree/profiling_cpu_events.md

@@ -63,7 +63,7 @@ subsequent commands analyzing the profile. Example:
 ```shell
 perf record -o /tmp/perf.data \
   ./tools/iree-benchmark-module \
-    --driver=local-task \
+    --device=local-task \
     ... command-line arguments of iree-benchmark-module as usual ...
 ```
 
@@ -73,7 +73,7 @@ by, with the `-e` flag. For instance, to sample by L1 cache misses, one may do:
 ```shell
 perf record -o /tmp/perf.data -e L1-dcache-load-misses \
   ./tools/iree-benchmark-module \
-    --driver=local-task \
+    --device=local-task \
     ... command-line arguments of iree-benchmark-module as usual ...
 ```
 
@@ -150,7 +150,7 @@ First, we record on the device:
 adb shell \
   simpleperf record -e raw-l1d-cache-refill -o /data/local/tmp/perf.data \
     /data/local/tmp/iree-benchmark-module \
-      --driver=local-task \
+      --device=local-task \
       ... command-line arguments of iree-benchmark-module as usual ...
 ```
 

docs/developers/developing_iree/profiling_vulkan_gpu.md

@@ -45,7 +45,7 @@ $ /path/to/iree/build/tools/iree-compile -- \
 # Then package the Android app
 $ /path/to/iree/source/tools/android/run_module_app/build_apk.sh \
   ./build-apk \
-  --driver vulkan \
+  --device vulkan \
   --module_file /tmp/mhlo-dot.vmfb \
   --entry_function dot \
   --function_input=...

docs/developers/developing_iree/profiling_with_tracy.md

@@ -126,7 +126,7 @@ Example:
 
 ```shell
 TRACY_NO_EXIT=1 /data/local/tmp/iree-benchmark-module \
-  --driver=local-task \
+  --device=local-task \
   --module_file=/data/local/tmp/android_module.fbvm \
   --entry_function=serving_default \
   --function_input=1x384xi32 \

docs/website/docs/building-from-source/android.md

@@ -127,7 +127,7 @@ adb push /tmp/simple_abs_vmvx.vmfb /data/local/tmp/
 Run the tool:
 
 ``` shell
-adb shell /data/local/tmp/iree-run-module --driver=local-task \
+adb shell /data/local/tmp/iree-run-module --device=local-task \
   --module_file=/data/local/tmp/simple_abs_vmvx.vmfb \
   --entry_function=abs \
   --function_input="f32=-5"

docs/website/docs/building-from-source/riscv.md

@@ -117,7 +117,7 @@ ${QEMU_BIN} \
   -cpu rv64 \
   -L ${RISCV_TOOLCHAIN_ROOT}/sysroot/ \
   ../iree-build-riscv/tools/iree-run-module \
-  --driver=local-task \
+  --device=local-task \
   --module_file=/tmp/simple_abs_vmvx.vmfb \
   --entry_function=abs \
   --function_input=f32=-5
@@ -163,7 +163,7 @@ ${QEMU_BIN} \
   -cpu rv64,x-v=true,x-k=true,vlen=256,elen=64,vext_spec=v1.0 \
   -L ${RISCV_TOOLCHAIN_ROOT}/sysroot/ \
   ../iree-build-riscv/tools/iree-run-module \
-  --driver=local-task \
+  --device=local-task \
   --module_file=mobilenet_cpu.vmfb \
   --entry_function=predict \
   --function_input="1x224x224x3xf32=0"

docs/website/docs/deployment-configurations/cpu.md

@@ -126,7 +126,7 @@ In the build directory, run the following command:
 
 ``` shell hl_lines="2"
 tools/iree-run-module \
-    --driver=local-task \
+    --device=local-task \
     --module_file=mobilenet_cpu.vmfb \
     --entry_function=predict \
     --function_input="1x224x224x3xf32=0"

docs/website/docs/deployment-configurations/gpu-cuda-rocm.md

@@ -158,7 +158,7 @@ In the build directory, run the following command:
 
     ``` shell hl_lines="2"
     tools/iree-run-module \
-        --driver=cuda \
+        --device=cuda \
         --module_file=mobilenet-cuda.vmfb \
         --entry_function=predict \
         --function_input="1x224x224x3xf32=0"
@@ -168,7 +168,7 @@ In the build directory, run the following command:
 
     ``` shell hl_lines="2"
     tools/iree-run-module \
-        --driver=rocm \
+        --device=rocm \
         --module_file=mobilenet-rocm.vmfb \
         --entry_function=predict \
         --function_input="1x224x224x3xf32=0"

docs/website/docs/deployment-configurations/gpu-vulkan.md

@@ -176,7 +176,7 @@ In the build directory, run the following command:
 
 ``` shell hl_lines="2"
 tools/iree-run-module \
-    --driver=vulkan \
+    --device=vulkan \
     --module_file=mobilenet-vulkan.vmfb \
     --entry_function=predict \
     --function_input="1x224x224x3xf32=0"

experimental/rocm/rocm_driver.c

@@ -135,9 +135,9 @@ static iree_status_t iree_hal_rocm_driver_query_available_devices(
     uint8_t* buffer_ptr =
         (uint8_t*)device_infos + device_count * sizeof(iree_hal_device_info_t);
     for (iree_host_size_t i = 0; i < device_count; ++i) {
-      hipDevice_t device;
-      iree_status_t status = ROCM_RESULT_TO_STATUS(
-          &driver->syms, hipDeviceGet(&device, i), "hipDeviceGet");
+      hipDevice_t device = 0;
+      status = ROCM_RESULT_TO_STATUS(&driver->syms, hipDeviceGet(&device, i),
+                                     "hipDeviceGet");
       if (!iree_status_is_ok(status)) break;
       buffer_ptr = iree_hal_rocm_populate_device_info(
           device, &driver->syms, buffer_ptr, &device_infos[i]);
@@ -152,6 +152,18 @@ static iree_status_t iree_hal_rocm_driver_query_available_devices(
   return status;
 }
 
+static iree_status_t iree_hal_rocm_driver_dump_device_info(
+    iree_hal_driver_t* base_driver, iree_hal_device_id_t device_id,
+    iree_string_builder_t* builder) {
+  iree_hal_rocm_driver_t* driver = iree_hal_rocm_driver_cast(base_driver);
+  hipDevice_t device = (hipDevice_t)device_id;
+  if (!device) return iree_ok_status();
+  // TODO: dump detailed device info.
+  (void)driver;
+  (void)device;
+  return iree_ok_status();
+}
+
 static iree_status_t iree_hal_rocm_driver_select_default_device(
     iree_hal_rocm_dynamic_symbols_t* syms, int default_device_index,
     iree_allocator_t host_allocator, hipDevice_t* out_device) {
@@ -219,6 +231,7 @@ static iree_status_t iree_hal_rocm_driver_create_device_by_path(
 static const iree_hal_driver_vtable_t iree_hal_rocm_driver_vtable = {
     .destroy = iree_hal_rocm_driver_destroy,
     .query_available_devices = iree_hal_rocm_driver_query_available_devices,
+    .dump_device_info = iree_hal_rocm_driver_dump_device_info,
     .create_device_by_id = iree_hal_rocm_driver_create_device_by_id,
     .create_device_by_path = iree_hal_rocm_driver_create_device_by_path,
 };

experimental/web/run_native_benchmarks.sh

@@ -34,7 +34,7 @@ IREE_BENCHMARK_MODULE_PATH=iree-benchmark-module
 echo "Benchmarking DeepLabV3..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./deeplabv3_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x257x257x3xf32 \
@@ -44,7 +44,7 @@ echo ""
 echo "Benchmarking MobileSSD..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./mobile_ssd_v2_float_coco_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x320x320x3xf32 \
@@ -54,7 +54,7 @@ echo ""
 echo "Benchmarking PoseNet..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./posenet_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x353x257x3xf32 \
@@ -64,7 +64,7 @@ echo ""
 echo "Benchmarking MobileBertSquad..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./mobilebertsquad_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x384xi32 \
@@ -76,7 +76,7 @@ echo ""
 echo "Benchmarking MobileNetV2..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./mobilenet_v2_1.0_224_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x224x224x3xf32 \
@@ -86,7 +86,7 @@ echo ""
 echo "Benchmarking MobileNetV3Small..."
 "${IREE_BENCHMARK_MODULE_PATH?}" \
     --module_file=./MobileNetV3SmallStaticBatch_native.vmfb \
-    --driver=local-task \
+    --device=local-task \
     --task_topology_group_count=1 \
     --entry_function=main \
     --function_input=1x224x224x3xf32 \

experimental/web/testing/parse_test_list.py

@@ -73,7 +73,7 @@ def parse_ctest_dump(ctest_dump_path, build_dir):
 
       # Parse the 'command' list into the source file and its arguments.
       #   /path/to/test_runner.js  # such as iree-check-module.js or test.js
-      #   arg 1                    # such as --driver=local-task
+      #   arg 1                    # such as --device=local-task
       #   arg 2                    # such as check_vmvx_op.mlir_module.vmfb
       test_source_absolute_path = test["command"][0]
       parsed_test["sourceFile"] = get_normalized_relative_path(

integrations/tensorflow/python_projects/iree_tf/iree/tf/support/trace_utils.py

@@ -285,7 +285,7 @@ def serialize(self, trace_dir: str) -> None:
         serialized_inputs = self.calls[0].serialized_inputs
         flagfile = [
             f"--module_file={compiled_path}",
-            f"--driver={self.backend_driver}",
+            f"--device={self.backend_driver}",
             f"--entry_function={entry_function}",
         ] + [f"--function_input={input}" for input in serialized_inputs]
         with open(os.path.join(trace_dir, "flagfile"), "w") as f: