[XLA] Add should_inline callback to CallInliner

Add a callback to CallInliner to allow finer-grained control over which callsites to inline.

PiperOrigin-RevId: 779849022

Author: mkuperst

third_party/xla/xla/service/BUILD

@@ -981,6 +981,7 @@ xla_cc_test(
     size = "small",
     srcs = ["call_inliner_test.cc"],
     deps = [
+        ":call_graph",
         ":call_inliner",
         "//xla:literal_util",
         "//xla:shape_util",

third_party/xla/xla/service/call_inliner.cc

@@ -298,6 +298,27 @@ bool CallInliner::IsInlineableCallOp(HloInstruction* instruction) const {
          InlineComposites(instruction, composites_to_preserve_);
 }
 
+bool CallInliner::ShouldInline(const CallGraph& call_graph,
+                               HloInstruction* instruction) const {
+  if (!IsInlineableCallOp(instruction)) {
+    return false;
+  }
+
+  if (should_inline_.has_value()) {
+    if (!(*should_inline_)(call_graph, instruction)) {
+      return false;
+    }
+  }
+
+  if (single_call_site_) {
+    return call_graph.GetNode(instruction->to_apply())
+               .caller_callsites()
+               .size() == 1;
+  }
+
+  return true;
+}
+
 absl::StatusOr<bool> CallInliner::InlineAndLegalize(
     const CallGraph& call_graph, HloComputation* computation,
     absl::Span<HloInstruction* const> instruction_sequence) const {
@@ -309,10 +330,7 @@ absl::StatusOr<bool> CallInliner::InlineAndLegalize(
     // used for parallel device computation.
     // TODO(b/229887502): update the inliner to ignore only parallel
     // device type async call instead of all.
-    if (IsInlineableCallOp(instruction) &&
-        (!single_call_site_ || call_graph.GetNode(instruction->to_apply())
-                                       .caller_callsites()
-                                       .size() == 1)) {
+    if (ShouldInline(call_graph, instruction)) {
       // The caller instruction will get removed after inlining. Record the
       // callee computation beforehand, so we can find its schedule.
       HloComputation* callee = instruction->to_apply();

third_party/xla/xla/service/call_inliner.h

@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef XLA_SERVICE_CALL_INLINER_H_
 #define XLA_SERVICE_CALL_INLINER_H_
 
+#include <functional>
+#include <optional>
 #include <string>
 #include <utility>
 
@@ -48,14 +50,19 @@ class CallInliner : public HloModulePass {
   // are being inlined if necessary.
   // If `uniquify_channel_ids` is true, the channel ids of the resulting
   // computation will be uniquified.
+  // If the callback `should_inline` is provided, only functions callsite for
+  // which it retuns true will be inlined.
   explicit CallInliner(
       bool single_call_site = false, bool update_domain = false,
       absl::flat_hash_set<std::string> composites_to_preserve = {},
-      bool uniquify_channel_ids = false)
+      bool uniquify_channel_ids = false,
+      std::optional<std::function<bool(const CallGraph&, HloInstruction*)>>
+          should_inline = std::nullopt)
       : single_call_site_(single_call_site),
         update_domain_(update_domain),
         uniquify_channel_ids_(uniquify_channel_ids),
-        composites_to_preserve_(std::move(composites_to_preserve)) {}
+        composites_to_preserve_(std::move(composites_to_preserve)),
+        should_inline_(std::move(should_inline)) {}
   ~CallInliner() override = default;
   absl::string_view name() const override { return "call-inliner"; }
 
@@ -73,10 +80,16 @@ class CallInliner : public HloModulePass {
       const CallGraph& call_graph, HloComputation* computation,
       absl::Span<HloInstruction* const> instruction_sequence) const;
 
+  bool ShouldInline(const CallGraph& call_graph,
+                    HloInstruction* instruction) const;
+
   bool single_call_site_;
   bool update_domain_;
   bool uniquify_channel_ids_;
   absl::flat_hash_set<std::string> composites_to_preserve_;
+  std::optional<
+      std::function<bool(const CallGraph& call_graph, HloInstruction*)>>
+      should_inline_;
 };
 
 }  // namespace xla

third_party/xla/xla/service/call_inliner_test.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "xla/literal_util.h"
+#include "xla/service/call_graph.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
@@ -813,5 +814,41 @@ ENTRY main {
   EXPECT_EQ(root->metadata().op_name(), "x/y");
 }
 
+TEST_F(CallInlinerTest, InliningCallBack) {
+  const char* hlo = R"(
+callee_negate {
+  input = f32[128,32] parameter(0)
+  ROOT y = f32[128,32] negate(input)
+}
+
+callee_trivial {
+  ROOT input = f32[128,32] parameter(0)
+}
+
+ENTRY main {
+  input = f32[128,32] parameter(0)
+  call.negate = f32[128,32] call(input), to_apply=callee_negate
+  call.trivial = f32[128,32] call(input), to_apply=callee_trivial
+  ROOT result = subtract(call.negate, call.trivial)
+})";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
+                          ParseAndReturnVerifiedModule(hlo));
+
+  auto inline_trivial_only = [](const CallGraph& call_graph,
+                                HloInstruction* instruction) {
+    HloComputation* callee = instruction->to_apply();
+    return (callee->root_instruction()->opcode() == HloOpcode::kParameter);
+  };
+  CallInliner call_inliner(/*single_call_site=*/false, /*update_domain=*/false,
+                           /*composites_to_preserve=*/{},
+                           /*uniquify_channel_ids=*/false,
+                           /*should_inline=*/inline_trivial_only);
+
+  EXPECT_THAT(call_inliner.Run(m.get()), ::tsl::testing::IsOkAndHolds(true));
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              op::Subtract(op::Call(op::Parameter(0)), op::Parameter(0)));
+}
+
 }  // namespace
 }  // namespace xla