Repository: iree-org/iree (language: python)

Issue:
Fix device list of loaded executable in PJRT plugin for multiple GPUs
It closes #19366, and blocks #19279.

After this PR, `ClientOptions::Compile` will first check the device assignment in the compile options, and then return the corresponding device list with the loaded executable.

To achieve this, we introduce protobuf via `FetchContent` in this PR, which is scoped to the PJRT plugin. Compile options will be passed by the PJRT client encoded in protobuf, and in this plugin we decode it first and then retrieve some interesting fields.

ci-exactly: build_packages, test_pjrt


Discussion / hints:
It works now and is ready for review : )

Fix patch:
diff --git a/integrations/pjrt/src/iree_pjrt/common/CMakeLists.txt b/integrations/pjrt/src/iree_pjrt/common/CMakeLists.txt
index 4371cd3db965..c5855fc76196 100644
--- a/integrations/pjrt/src/iree_pjrt/common/CMakeLists.txt
+++ b/integrations/pjrt/src/iree_pjrt/common/CMakeLists.txt
@@ -56,6 +56,7 @@ iree_cc_library(
       iree::compiler::bindings::c::loader
       iree_pjrt::partitioner_api
       iree_pjrt::partitioner_api::loader
+      iree_pjrt_deps::protos
     PUBLIC
 )
 
diff --git a/integrations/pjrt/src/iree_pjrt/common/api_impl.cc b/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
index 68b45ecbbb07..6f2d787ea1d2 100644
--- a/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
+++ b/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
@@ -1487,8 +1487,8 @@ PJRT_Error* ClientInstance::Compile(const PJRT_Program* program,
     }
 
     // Set flags.
-    // TODO: Plumb CompileOptions through.
-    // if (!job->SetFlags(options)) return MakeCompilerError(*job);
+    if (!job->SetFlags(options)) return MakeCompilerError(*job);
+
     if (artifact_tx) {
       artifact_tx->WriteArtifact(
           /*label=*/"partitioner_flags", /*extension=*/"txt", /*index=*/-1,
@@ -1538,8 +1538,8 @@ PJRT_Error* ClientInstance::Compile(const PJRT_Program* program,
     if (!SetDefaultCompilerFlags(job.get())) {
       return MakeCompilerError(*job);
     }
-    // TODO: Plumb CompileOptions through.
-    // if (!job->SetFlags(options)) return MakeCompilerError(*job);
+    if (!job->SetFlags(options)) return MakeCompilerError(*job);
+
     if (artifact_tx) {
       artifact_tx->WriteArtifact(
           /*label=*/"flags", /*extension=*/"txt", /*index=*/-1,
diff --git a/integrations/pjrt/src/iree_pjrt/common/compiler.h b/integrations/pjrt/src/iree_pjrt/common/compiler.h
index faf07dbcb318..9969a00db8dc 100644
--- a/integrations/pjrt/src/iree_pjrt/common/compiler.h
+++ b/integrations/pjrt/src/iree_pjrt/common/compiler.h
@@ -11,8 +11,7 @@
 #include <string>
 
 #include "iree_pjrt/common/debugging.h"
-// TODO: Excise.
-// #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/compile_options.pb.h"
 
 namespace iree::pjrt {
 
@@ -37,8 +36,7 @@ class CompilerJob {
   // setup of a job (or if the underlying session will not be re-used).
   // Returns false on failure.
   virtual bool SetFlag(const char* flag) = 0;
-  // TODO: Excise.
-  // virtual bool SetFlags(xla::CompileOptions options) = 0;
+  virtual bool SetFlags(xla::CompileOptionsProto options) = 0;
 
   // Gets all flags as a string. This is intended for debug printing a plausible
   // command line to reproduce compilation.
diff --git a/integrations/pjrt/src/iree_pjrt/common/hlo_partitioner.cc b/integrations/pjrt/src/iree_pjrt/common/hlo_partitioner.cc
index b3238076d927..c00f6eb0c88b 100644
--- a/integrations/pjrt/src/iree_pjrt/common/hlo_partitioner.cc
+++ b/integrations/pjrt/src/iree_pjrt/common/hlo_partitioner.cc
@@ -97,41 +97,44 @@ class OpenXLAPartitionerJob : public CompilerJob {
     return true;
   }
 
-  // TODO: Find another way to deal with this.
-  // bool SetFlags(xla::CompileOptions options) override {
-  //   int num_partitions = options.executable_build_options.num_partitions();
-  //   int num_replicas = options.executable_build_options.num_replicas();
-  //   bool use_spmd_partitioning =
-  //       options.executable_build_options.use_spmd_partitioning();
-  //   auto allow_spmd_sharding_propagation_to_output =
-  //       options.executable_build_options
-  //           .allow_spmd_sharding_propagation_to_output();
-  //   if (!SetFlag(absl::StrCat("--openxla-partitioner-gspmd-num-partitions=",
-  //                             num_partitions)
-  //                    .c_str())) {
-  //     return false;
-  //   }
-  //   if (!SetFlag(absl::StrCat("--openxla-partitioner-gspmd-replica-count=",
-  //                             num_replicas)
-  //                    .c_str())) {
-  //     return false;
-  //   }
-  //   if (!SetFlag(
-  //           absl::StrCat("--openxla-partitioner-gspmd-use-spmd-partitioning=",
-  //                        use_spmd_partitioning)
-  //               .c_str())) {
-  //     return false;
-  //   }
-  //   if (!SetFlag(
-  //           absl::StrCat(
-  //               "--openxla-partitioner-gspmd-allow-spmd-"
-  //               "sharding-propagation-to-output=",
-  //               absl::StrJoin(allow_spmd_sharding_propagation_to_output,
-  //               ",")) .c_str())) {
-  //     return false;
-  //   }
-  //   return true;
-  // }
+  bool SetFlags(xla::CompileOptionsProto options) override {
+    int num_partitions = options.executable_build_options().num_partitions();
+    int num_replicas = options.executable_build_options().num_replicas();
+    bool use_spmd_partitioning =
+        options.executable_build_options().use_spmd_partitioning();
+    auto allow_spmd_sharding_propagation_to_output =
+        options.executable_build_options()
+            .allow_spmd_sharding_propagation_to_output();
+    if (!SetFlag(("--openxla-partitioner-gspmd-num-partitions=" +
+                  std::to_string(num_partitions))
+                     .c_str())) {
+      return false;
+    }
+    if (!SetFlag(("--openxla-partitioner-gspmd-replica-count=" +
+                  std::to_string(num_replicas))
+                     .c_str())) {
+      return false;
+    }
+    if (!SetFlag(("--openxla-partitioner-gspmd-use-spmd-partitioning=" +
+                  std::to_string(use_spmd_partitioning))
+                     .c_str())) {
+      return false;
+    }
+    std::string allow_spmd_sharding_propagation_to_output_str;
+    for (size_t i = 0; i < allow_spmd_sharding_propagation_to_output.size();
+         ++i) {
+      if (i != 0) allow_spmd_sharding_propagation_to_output_str += ",";
+      allow_spmd_sharding_propagation_to_output_str +=
+          std::to_string(allow_spmd_sharding_propagation_to_output[i]);
+    }
+    if (!SetFlag(("--openxla-partitioner-gspmd-allow-spmd-"
+                  "sharding-propagation-to-output=" +
+                  allow_spmd_sharding_propagation_to_output_str)
+                     .c_str())) {
+      return false;
+    }
+    return true;
+  }
 
   std::string GetFlags() override {
     std::string flags;
diff --git a/integrations/pjrt/src/iree_pjrt/common/iree_compiler.cc b/integrations/pjrt/src/iree_pjrt/common/iree_compiler.cc
index 095e540fa3a0..1cddbc3b37b6 100644
--- a/integrations/pjrt/src/iree_pjrt/common/iree_compiler.cc
+++ b/integrations/pjrt/src/iree_pjrt/common/iree_compiler.cc
@@ -95,29 +95,28 @@ class IREECompilerJob : public CompilerJob {
     return true;
   }
 
-  // TODO: Excise: Cannot dep on an internal XLA structure.
-  // bool SetFlags(xla::CompileOptions options) override {
-  //   // Set extra options, overriding env variables if appropriate.
-  //   for (auto [option, option_override] : options.env_option_overrides) {
-  //     std::string override_string;
-  //     if (auto override_val = std::get_if<std::string>(&option_override)) {
-  //       override_string = *override_val;
-  //     } else if (auto override_val = std::get_if<bool>(&option_override)) {
-  //       override_string = *override_val ? "true" : "false";
-  //     } else if (auto override_val = std::get_if<int64_t>(&option_override))
-  //     {
-  //       override_string = std::to_string(*override_val);
-  //     } else {
-  //       assert(false &&
-  //              "option value should be of type string, bool, or int64");
-  //     }
-  //     if (!SetFlag(absl::StrCat("--", option, "=", override_string).c_str()))
-  //     {
-  //       return false;
-  //     }
-  //   }
-  //   return true;
-  // }
+  bool SetFlags(xla::CompileOptionsProto options) override {
+    // Set extra options, overriding env variables if appropriate.
+    for (auto [option, option_override] : options.env_option_overrides()) {
+      std::string override_string;
+      if (option_override.has_string_field()) {
+        override_string = option_override.string_field();
+      } else if (option_override.has_bool_field()) {
+        override_string = option_override.bool_field() ? "true" : "false";
+      } else if (option_override.has_int_field()) {
+        override_string = std::to_string(option_override.int_field());
+      } else if (option_override.has_double_field()) {
+        override_string = std::to_string(option_override.double_field());
+      } else {
+        assert(false &&
+               "option value should be of type string, bool, int, or double");
+      }
+      if (!SetFlag(("--" + option + "=" + override_string).c_str())) {
+        return false;
+      }
+    }
+    return true;
+  }
 
   std::string GetFlags() override {
     std::string flags;


Test patch:
diff --git a/.github/workflows/pkgci_test_pjrt.yml b/.github/workflows/pkgci_test_pjrt.yml
index a32b1c4503da..cc5fba10e386 100644
--- a/.github/workflows/pkgci_test_pjrt.yml
+++ b/.github/workflows/pkgci_test_pjrt.yml
@@ -61,7 +61,7 @@ jobs:
           source ${VENV_DIR}/bin/activate
           python -m pip install -v --no-deps -e integrations/pjrt/python_packages/iree_${{ matrix.pjrt_platform }}_plugin
           # install
-          python -m pip install jax==0.4.35
+          python -m pip install jax==0.4.36
       - name: Run tests
         run: |
           source ${VENV_DIR}/bin/activate
diff --git a/build_tools/testing/run_jax_tests.sh b/build_tools/testing/run_jax_tests.sh
index 66b21351d8cc..a065d16890e9 100755
--- a/build_tools/testing/run_jax_tests.sh
+++ b/build_tools/testing/run_jax_tests.sh
@@ -48,6 +48,13 @@ diff_jax_test test/test_add.py
 diff_jax_test test/test_degenerate.py
 diff_jax_test test/test_simple.py
 
+# here we test if the compile options is passed to IREE PJRT plugin successfully.
+# we pass --iree-scheduling-dump-statistics-format=csv via jax.jit,
+# and see if there's statistics in the output
+compile_options_test_tmp_out=$(mktemp /tmp/jax_test_result_compile_options.XXXXXX)
+JAX_PLATFORMS=$actual_jax_platform python test/test_compile_options.py 2>&1 | tee $compile_options_test_tmp_out
+cat $compile_options_test_tmp_out | grep '@main_dispatch'
+
 
 # FIXME: we can also utilize the native test cases from JAX,
 # e.g. `tests/nn_test.py` from the JAX repo, as below,
diff --git a/integrations/pjrt/test/test_compile_options.py b/integrations/pjrt/test/test_compile_options.py
new file mode 100644
index 000000000000..877ab9911ed7
--- /dev/null
+++ b/integrations/pjrt/test/test_compile_options.py
@@ -0,0 +1,19 @@
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from functools import partial
+import jax.numpy as jnp
+from jax import jit
+
+a = jnp.asarray([1, 2, 3, 4, 5, 6, 7, 8, 9])
+
+
+@partial(jit, compiler_options={"iree-scheduling-dump-statistics-format": "csv"})
+def f(a, b):
+    return a + b
+
+
+print(f(a, a))