Create a PjRt interpreter client.

The existing stream executor interpreter platform is used throughout our testing frameworks as a reference platform. It is implemented as a wrapper around `HloEvaluator`. This change adds a new PjRt client to act as an alternative to the stream executor interpreter platform. It wraps `HloEvaluator`, like the stream executor based implementation. This is a first implementation that leaves many interfaces stubbed instead of implemented. It has been validated against all tests that have been migrated to `HloPjRtTestBase`. PiperOrigin-RevId: 706070034
openxla · Dec 19, 2024 · 7b2617b · 7b2617b
1 parent 43a5ba5
commit 7b2617b
Show file tree

Hide file tree

Showing 3 changed files with 1,060 additions and 0 deletions.
diff --git a/xla/pjrt/interpreter/BUILD b/xla/pjrt/interpreter/BUILD
@@ -0,0 +1,67 @@
+load("//xla/tsl:tsl.bzl", "internal_visibility")
+load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "interpreter_client",
+    srcs = ["interpreter_client.cc"],
+    hdrs = ["interpreter_client.h"],
+    visibility = internal_visibility(["//xla:friends"]),
+    deps = [
+        "//xla:literal",
+        "//xla:shape_util",
+        "//xla:util",
+        "//xla/backends/interpreter:compiler",
+        "//xla/client:executable_build_options",
+        "//xla/hlo/builder:xla_computation",
+        "//xla/hlo/evaluator:hlo_evaluator",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass_pipeline",
+        "//xla/hlo/transforms:cholesky_expander",
+        "//xla/hlo/transforms:dynamic_index_splitter",
+        "//xla/hlo/transforms:eigh_expander",
+        "//xla/hlo/transforms:qr_expander",
+        "//xla/pjrt:layout_mode",
+        "//xla/pjrt:mlir_to_hlo",
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt:pjrt_common",
+        "//xla/pjrt:pjrt_compiler",
+        "//xla/pjrt:pjrt_device_description",
+        "//xla/pjrt:pjrt_executable",
+        "//xla/pjrt:pjrt_future",
+        "//xla/pjrt:utils",
+        "//xla/service:batchnorm_expander",
+        "//xla/service:computation_placer_hdr",
+        "//xla/service:custom_call_target_registry",
+        "//xla/service:dynamic_dimension_inference",
+        "//xla/service:hlo_cost_analysis",
+        "//xla/service:hlo_module_config",
+        "//xla/service:hlo_module_util",
+        "//xla/service:layout_assignment",
+        "//xla/service:topk_rewriter",
+        "//xla/service:triangular_solve_expander",
+        "//xla/tsl/platform:errors",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:die_if_null",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
+        "@tsl//tsl/platform:fingerprint",
+        "@tsl//tsl/platform:statusor",
+    ],
+)