Adding test for iree-run-module with multiple devices.

tools/test/BUILD.bazel

diff --git a/tools/test/BUILD.bazel b/tools/test/BUILD.bazel
index cf5878c..46709aa 100644
--- a/tools/test/BUILD.bazel
+++ b/tools/test/BUILD.bazel
@@ -31,6 +31,7 @@
             "iree-run-mlir.mlir",
             "iree-run-module-expected.mlir",
             "iree-run-module-inputs.mlir",
+            "iree-run-module-multi.mlir",
             "iree-run-module-outputs.mlir",
             "iree-run-module.mlir",
             "multiple_args.mlir",

tools/test/CMakeLists.txt

diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt
index 75dde66..a866548 100644
--- a/tools/test/CMakeLists.txt
+++ b/tools/test/CMakeLists.txt
@@ -27,6 +27,7 @@
     "iree-run-mlir.mlir"
     "iree-run-module-expected.mlir"
     "iree-run-module-inputs.mlir"
+    "iree-run-module-multi.mlir"
     "iree-run-module-outputs.mlir"
     "iree-run-module.mlir"
     "multiple_args.mlir"

tools/test/iree-run-module-multi.mlir

diff --git a/tools/test/iree-run-module-multi.mlir b/tools/test/iree-run-module-multi.mlir
new file mode 100644
index 0000000..3412596
--- /dev/null
+++ b/tools/test/iree-run-module-multi.mlir
@@ -0,0 +1,43 @@
+// Tests that multiple devices are supported through iree-run-module by
+// providing two local thread pools. This is not optimal and not an intended
+// route for multi-device CPU workloads but requires no additional hardware
+// resources for the test and still verifies the compiler/runtime tooling
+// rendezvous of devices as specified on the command line.
+
+// RUN: (iree-compile %s \
+// RUN:      --iree-execution-model=async-external \
+// RUN:      --iree-hal-target-device=device_a=local[0] \
+// RUN:      --iree-hal-target-device=device_b=local[1] \
+// RUN:      --iree-hal-local-target-device-backends=vmvx | \
+// RUN:  iree-run-module \
+// RUN:      --module=- \
+// RUN:      --function=mutli_device_mul \
+// RUN:      --input=4xf32=10,11,12,13 \
+// RUN:      --device=local-task \
+// RUN:      --device=local-task \
+// RUN:      --task_topology_group_count=1) | \
+// RUN: FileCheck %s
+
+// CHECK: EXEC @mutli_device_mul
+// CHECK-NEXT: result[0]: hal.buffer_view
+// CHECK-NEXT: 4xf32=0 55 144 273
+func.func public @mutli_device_mul(
+  // Input argument is resident on device_a (tooling default to first device).
+  %input_a: tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>}
+) -> (
+  // Output result is expected to be on device_a (though not required).
+  tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>}
+) {
+  // Compute on device_a (input is there).
+  %constant_a = arith.constant dense<[0.0, 1.0, 2.0, 3.0]> : tensor<4xf32>
+  %transient_a = arith.mulf %input_a, %constant_a : tensor<4xf32>
+  // Transfer the result from device_a -> device_b.
+  %transient_b = flow.tensor.transfer %transient_a : tensor<4xf32> to #hal.device.promise<@device_b>
+  // Compute on device_b.
+  %constant_b = arith.constant dense<[4.0, 5.0, 6.0, 7.0]> : tensor<4xf32>
+  %result_b = arith.mulf %transient_b, %constant_b : tensor<4xf32>
+  // Transfer the result from device_b -> device_a.
+  %result_a = flow.tensor.transfer %result_b : tensor<4xf32> to #hal.device.promise<@device_a>
+  // Return the result on device_a (as required by ABI attr).
+  func.return %result_a : tensor<4xf32>
+}