[MPMD-GPU] Add is_subslice_topology to the IFRT Topology.

This is to make IFRT Topology fields consistent with `PjRtTopologyDescriptionProto` fields.

PiperOrigin-RevId: 694647266

Author: changhuilin

xla/debug_options_flags.cc

@@ -2050,6 +2050,16 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
                 bool_setter_for(&DebugOptions::set_xla_enable_fast_math),
                 debug_options->xla_enable_fast_math(),
                 "Enable optimizations that assume finite math, i.e., no NaN."));
+  flag_list->push_back(tsl::Flag(
+      "xla_experimental_exec_time_optimization_effort",
+      float_setter_for(
+          &DebugOptions::set_xla_experimental_exec_time_optimization_effort),
+      debug_options->xla_experimental_exec_time_optimization_effort(),
+      "The execution time optimization effort to expend during compilation. "
+      "Takes range [-1.0, 1.0] where values < 0.0 indicate skipping passes "
+      "which might optimize the final runtime (thus improving compile time), "
+      "and values > 0.0 indicate running additional passes which may improve "
+      "runtime at the cost of compilation time."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_experimental_parallel_collective_overlap_limit",
       int32_setter_for(

xla/python/ifrt/topology.h

@@ -44,6 +44,8 @@ class Topology : public llvm::RTTIExtends<Topology, llvm::RTTIRoot> {
 
   virtual PjRtPlatformId platform_id() const = 0;
 
+  virtual bool is_subslice_topology() const = 0;
+
   // Returns an unordered list of descriptions for all devices in this topology.
   // TODO(phawkins): consider introducing an IFRT-specific API here instead of
   // delegating to PJRT.

xla/python/pjrt_ifrt/pjrt_topology.cc

@@ -51,6 +51,10 @@ PjRtPlatformId PjRtTopology::platform_id() const {
   return description_->platform_id();
 }
 
+bool PjRtTopology::is_subslice_topology() const {
+  return description_->is_subslice_topology();
+}
+
 std::vector<std::unique_ptr<const PjRtDeviceDescription>>
 PjRtTopology::DeviceDescriptions() const {
   return description_->DeviceDescriptions();

xla/python/pjrt_ifrt/pjrt_topology.h

@@ -46,6 +46,7 @@ class PjRtTopology final : public llvm::RTTIExtends<PjRtTopology, Topology> {
   absl::string_view platform_name() const override;
   absl::string_view platform_version() const override;
   PjRtPlatformId platform_id() const override;
+  bool is_subslice_topology() const override;
 
   std::vector<std::unique_ptr<const PjRtDeviceDescription>> DeviceDescriptions()
       const override;

xla/python/xla.cc

@@ -858,6 +858,13 @@ NB_MODULE(xla_extension, m_nb) {
       .def_prop_ro(
           "platform_version",
           [](ifrt::Topology& topology) { return topology.platform_version(); })
+      .def_prop_ro(
+          "platform_id",
+          [](ifrt::Topology& topology) { return topology.platform_id(); })
+      .def_prop_ro("is_subslice_topology",
+                   [](ifrt::Topology& topology) {
+                     return topology.is_subslice_topology();
+                   })
       .def("serialize",
            [](ifrt::Topology& topology) -> nb::bytes {
              std::string serialized = ValueOrThrow(topology.Serialize());

xla/python/xla_extension/__init__.pyi

@@ -738,6 +738,8 @@ class Executable:
 class DeviceTopology:
   platform: str
   platform_version: str
+  platform_id: int
+  is_subslice_topology: bool
   def _make_compile_only_devices(self) -> List[Device]: ...
   def serialize(self) -> bytes: ...
   def __getattr__(self, name: str) -> Any: ...

xla/service/gpu/transforms/softmax_rewriter_triton.cc

@@ -124,109 +124,6 @@ inline bool HasOneUse(const HloInstruction* instr) {
   return instr->user_count() == 1;
 }
 
-// Supports two types of broadcast of parameters. Either to one batch
-// dim, or one reduction dim. For example the following cases are supported:
-//
-// Case #1:
-// p = f32[a] parameter(0)
-// b = f32[a,x] broadcast(p), dimensions={0}
-//
-// Case #2:
-// p = f32[a] parameter(0)
-// b = f32[x,a] broadcast(p), dimensions={1}
-//
-// Case #3:
-// p = f32[a,b] parameter(0)
-// b = f32[x,a,b] broadcast(p), dimensions={1,2}
-//
-// Other broadcast tiling patterns are currently unsupported.
-// See b/328049138 for details.
-//
-// Unsupported case #1:
-// p = f32[a] parameter(0)
-// b = f32[x,a,y] broadcast(p), dimensions={1}
-//
-// Unsupported case #2:
-// p = f32[a,b] parameter(0)
-// b = f32[x,a,y,b] broadcast(p), dimensions={1,3}
-//
-// Unsupported case #3:
-// p = f32[a] parameter(0)
-// b = f32[x,y,a] broadcast(p), dimensions={2}
-//
-// Unsupported case #4:
-// p = f32[a,b] parameter(0)
-// b = f32[a,x,b] broadcast(p), dimensions={0,2}
-//
-// Unsupported case #5:
-// p = f32[] parameter(0)
-// b = f32[x] broadcast(p), dimensions={}
-bool IsBatchOrReductionDimBroadcast(const HloInstruction& hlo) {
-  CHECK_EQ(hlo.opcode(), HloOpcode::kBroadcast)
-      << "Expected broadcast " << hlo.ToShortString();
-  CHECK_EQ(hlo.operand(0)->opcode(), HloOpcode::kParameter)
-      << "Expected parameter " << hlo.operand(0)->ToShortString();
-
-  const HloBroadcastInstruction* broadcast =
-      Cast<HloBroadcastInstruction>(&hlo);
-
-  const HloParameterInstruction* parameter =
-      Cast<HloParameterInstruction>(hlo.operand(0));
-
-  // Support only one dim broadcast. Scalar parameters are handled elsewhere.
-  if (broadcast->dimensions().empty() ||
-      parameter->shape().dimensions_size() + 1 !=
-          broadcast->shape().dimensions_size()) {
-    return false;
-  }
-
-  // It is enough to ensure that the broadcast does not preserve both last, and
-  // first dimensions of the parameter at the same time. Otherwise the broadcast
-  // is the unsupported case #4.
-  //
-  // Preserve the first dim:
-  //   p = f32[a,b] parameter(0)
-  //   b1 = f32[a,b,c] broadcast(p), dimensions={0,1}
-  bool preserve_first_dim = broadcast->dimensions().front() == 0;
-  // Preserve the last dim:
-  //   p = f32[a,b] parameter(0)
-  //   b1 = f32[c,a,b] broadcast(p), dimensions={1,2}
-  bool preserve_last_dim = broadcast->dimensions().back() ==
-                           broadcast->shape().dimensions_size() - 1;
-  // We do not want to preserve both first and last dim, as it means the
-  // broadcast is not expanding on outermost dims.
-  return !(preserve_first_dim && preserve_last_dim);
-}
-
-bool IsBroadcastOfAScalar(const HloInstruction& hlo) {
-  CHECK_EQ(hlo.opcode(), HloOpcode::kBroadcast)
-      << "Expected broadcast " << hlo.ToShortString();
-  return ShapeUtil::IsScalar(hlo.operand(0)->shape());
-}
-
-bool IsSingleRowParameterBroadcast(const HloInstruction& hlo) {
-  CHECK_EQ(hlo.opcode(), HloOpcode::kBroadcast)
-      << "Expected broadcast " << hlo.ToShortString();
-  CHECK_EQ(hlo.operand(0)->opcode(), HloOpcode::kParameter)
-      << "Expected parameter " << hlo.operand(0)->ToShortString();
-
-  const HloBroadcastInstruction* broadcast =
-      Cast<HloBroadcastInstruction>(&hlo);
-  const HloParameterInstruction* parameter =
-      Cast<HloParameterInstruction>(hlo.operand(0));
-
-  if (parameter->shape().dimensions_size() != 1) {
-    return false;
-  }
-  return broadcast->dimensions()[0] == broadcast->shape().dimensions_size() - 1;
-}
-
-bool IsSupportedBroadcastOfParameter(const HloInstruction& hlo) {
-  return IsBroadcastOfParameter(hlo) &&
-         (IsBatchOrReductionDimBroadcast(hlo) || IsBroadcastOfAScalar(hlo) ||
-          IsSingleRowParameterBroadcast(hlo));
-}
-
 // Chooses which operand to use for fusion processing. Taking in a unary or
 // binary instruction, returns the first non-splat operand. If none is
 // present, returns any operand.
@@ -238,7 +135,7 @@ HloInstruction* ChooseOperandForFusionProcessing(HloInstruction* instr) {
   // broadcast of any op.
   if (instr->operand_count() > 1 &&
       (IsBroadcastOfScalarConstant(*instr->operand(0)) ||
-       IsSupportedBroadcastOfParameter(*instr->operand(0)))) {
+       IsBroadcastOfParameter(*instr->operand(0)))) {
     return instr->mutable_operand(1);
   }
   return instr->mutable_operand(0);
@@ -284,9 +181,9 @@ bool IsTriviallyFusible(HloInstruction* instr,
     // TODO(b/326217416): Extend the broadcast of splat constants/parameters to
     // a broadcast of any op.
     if ((IsBroadcastOfScalarConstant(*operand_0) ||
-         IsSupportedBroadcastOfParameter(*operand_0)) ^
+         IsBroadcastOfParameter(*operand_0)) ^
         (IsBroadcastOfScalarConstant(*operand_1) ||
-         IsSupportedBroadcastOfParameter(*operand_1))) {
+         IsBroadcastOfParameter(*operand_1))) {
       return static_cast<bool>(
           IsTritonSupportedInstruction(*instr, gpu_version));
     }

xla/service/gpu/transforms/softmax_rewriter_triton_test.cc

@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/transforms/softmax_rewriter_triton.h"
 
-#include <cstdint>
 #include <memory>
 #include <string>
 #include <variant>
@@ -1436,9 +1435,8 @@ ENTRY main {
   EXPECT_TRUE(fusion_rewriter_.Run(module.get()).value());
 }
 
-TEST_F(
-    SoftmaxRewriterTritonTest,
-    DoesNotFuseBinaryElementwiseIfIntermediateDiamondOpIsBroadcastOf1DParameterAlongBothBatchAndReductionDimensions) {  // NOLINT(whitespace/line_length)
+TEST_F(SoftmaxRewriterTritonTest,
+       FusesBinaryElementwiseIfIntermediateDiamondOpIsBroadcastOfParameter) {
   const std::string hlo_string = R"(
 HloModule h1
 
@@ -1460,67 +1458,12 @@ ENTRY main {
   ROOT add1 = f32[64,32,16]{2,1,0} add(add_0, broadcast_0)
 })";
   auto module = ParseAndReturnVerifiedModule(hlo_string).value();
-  EXPECT_FALSE(fusion_rewriter_.Run(module.get()).value());
-}
-
-TEST_F(
-    SoftmaxRewriterTritonTest,
-    DoesNotFuseBinaryElementwiseIfIntermediateDiamondOpWithBroadcastAlongBatchAndReductionDimAsParameter) {  // NOLINT(whitespace/line_length)
-  const std::string hlo_string = R"(
-HloModule h1
-
-add_computation {
-  y = f32[] parameter(1)
-  x = f32[] parameter(0)
-  ROOT add = f32[] add(x, y)
-}
-
-ENTRY main {
-  p0 = f32[8]{0} parameter(0)
-  p1 = f32[32,8,16]{2,1,0} parameter(1)
-  c = f32[] constant(0)
-
-  r0 = f32[32,8]{1,0} reduce(p1, c), dimensions={2}, to_apply=add_computation
-  b0 = f32[32,8,16]{2,1,0} broadcast(r0), dimensions={0,1}
-  b1 = f32[32,8,16]{2,1,0} broadcast(p0), dimensions={1}
-  add0 = f32[32,8,16]{2,1,0} add(b1, p1)
-  ROOT add1 = f32[32,8,16]{2,1,0} add(add0, b0)
-})";
-  auto module = ParseAndReturnVerifiedModule(hlo_string).value();
-  EXPECT_FALSE(fusion_rewriter_.Run(module.get()).value());
-}
-
-TEST_F(
-    SoftmaxRewriterTritonTest,
-    DoesNotFuseBinaryElementwiseIfIntermediateDiamondOpWithPartialBroadcastToBatchDim) {  // NOLINT(whitespace/line_length)
-  const std::string hlo_string = R"(
-HloModule h1
-
-add_computation {
-  y = f32[] parameter(1)
-  x = f32[] parameter(0)
-  ROOT add = f32[] add(x, y)
-}
-
-ENTRY main {
-  p0 = f32[16,64]{1,0} parameter(0)
-  p1 = f32[8,16,32,64]{3,2,1,0} parameter(1)
-  c = f32[] constant(0)
-
-  r0 = f32[8,16,32]{2,1,0} reduce(p1, c), dimensions={3}, to_apply=add_computation
-  b0 = f32[8,16,32,64]{3,2,1,0} broadcast(r0), dimensions={0,1,2}
-  b1 = f32[8,16,32,64]{3,2,1,0} broadcast(p0), dimensions={1,3}
-  add0 = f32[8,16,32,64]{3,2,1,0} add(b1, p1)
-  ROOT add1 = f32[8,16,32,64]{3,2,1,0} add(add0, b0)
-}
-)";
-  auto module = ParseAndReturnVerifiedModule(hlo_string).value();
-  EXPECT_FALSE(fusion_rewriter_.Run(module.get()).value());
+  EXPECT_TRUE(fusion_rewriter_.Run(module.get()).value());
 }
 
 TEST_F(
     SoftmaxRewriterTritonTest,
-    DoesNotFuseBinaryElementwiseIfIntermediateDiamondOpWithMultiDimBroadcastAlongBatchDimAsParameter) {  // NOLINT(whitespace/line_length)
+    FusesBinaryElementwiseIfIntermediateDiamondOpWithMultipleDimensionsAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
 HloModule h1
 
@@ -1542,7 +1485,7 @@ ENTRY main {
   ROOT add1 = f32[128,64,32,16]{3,2,1,0} add(add0, b0)
 })";
   auto module = ParseAndReturnVerifiedModule(hlo_string).value();
-  EXPECT_FALSE(fusion_rewriter_.Run(module.get()).value());
+  EXPECT_TRUE(fusion_rewriter_.Run(module.get()).value());
 }
 
 // Triton has a requirement that any tile in the program should not have more

xla/xla.proto

@@ -43,6 +43,14 @@ message DebugOptions {
   //--------------------------------------------------------------------------//
   // go/keep-sorted start
 
+  // The execution time optimization effort to expend during compilation.
+  // See `exec_time_optimization_effort` for accepted ranges. This flag will
+  // override any changes set in `ExecutionOptions`. Most likely this is just a
+  // temporary measure before https://github.com/jax-ml/jax/issues/24715 is in.
+  //
+  // TODO(b/377871215): Check whether we still need this.
+  float xla_experimental_exec_time_optimization_effort = 346;
+
   // go/keep-sorted end
 
   //--------------------------------------------------------------------------//
@@ -1051,7 +1059,7 @@ message DebugOptions {
   // be deterministic, although with additional overhead.
   bool xla_gpu_enable_scatter_determinism_expander = 345;
 
-  // Next id: 346
+  // Next id: 347
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.