tf-nightly-cpu 2.20.0.dev20250220__cp310-cp310-win_amd64.whl → 2.20.0.dev20250222__cp310-cp310-win_amd64.whl

tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_casting_utils.h

@@ -44,28 +44,6 @@ T* Cast(HloInstruction* instr) {
   return tsl::down_cast<T*>(instr);
 }
 
-// Downcasts a const HloInstruction pointer or returns nullptr if argument is
-// nullptr. Dies if TargetClass::ClassOf() does not match.
-template <typename T>
-const T* CastOrNull(const HloInstruction* i) {
-  if (i == nullptr) {
-    return nullptr;
-  }
-  CHECK(T::ClassOf(i));
-  return tsl::down_cast<const T*>(i);
-}
-
-// Downcasts a const HloInstruction pointer or returns nullptr if argument is
-// nullptr. Dies if TargetClass::ClassOf() does not match.
-template <typename T>
-T* CastOrNull(HloInstruction* i) {
-  if (i == nullptr) {
-    return nullptr;
-  }
-  CHECK(T::ClassOf(i));
-  return tsl::down_cast<T*>(i);
-}
-
 // Downcasts a const HloInstruction pointer or returns nullptr if
 // TargetClass::ClassOf() does not match. Dies if argument is nullptr. Similar
 // to LLVM's dyn_cast.
@@ -84,28 +62,6 @@ T* DynCast(HloInstruction* i) {
   return !T::ClassOf(i) ? nullptr : tsl::down_cast<T*>(i);
 }
 
-// Downcasts a const HloInstruction pointer. Return nullptr if argument is
-// nullptr orTargetClass::ClassOf() does not match. Similar to LLVM's
-// dyn_cast_or_null.
-template <typename T>
-const T* DynCastOrNull(const HloInstruction* instruction) {
-  if (instruction == nullptr || !T::ClassOf(instruction)) {
-    return nullptr;
-  }
-  return tsl::down_cast<const T*>(instruction);
-}
-
-// Downcasts a non-const HloInstruction pointer. Return nullptr if argument is
-// nullptr orTargetClass::ClassOf() does not match. Similar to LLVM's
-// dyn_cast_or_null.
-template <typename T>
-T* DynCastOrNull(HloInstruction* instruction) {
-  if (instruction == nullptr || !T::ClassOf(instruction)) {
-    return nullptr;
-  }
-  return tsl::down_cast<T*>(instruction);
-}
-
 }  // namespace xla
 
 #endif  // XLA_HLO_IR_HLO_CASTING_UTILS_H_

tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_instruction.h

@@ -1914,6 +1914,18 @@ class HloInstruction {
                           result_accuracy().mode() != ResultAccuracy::DEFAULT);
   }
 
+  bool equal_result_accuracy(const HloInstruction* other) const {
+    return result_accuracy().has_tolerance() ==
+               other->result_accuracy().has_tolerance() &&
+           result_accuracy().tolerance().atol() ==
+               other->result_accuracy().tolerance().atol() &&
+           result_accuracy().tolerance().rtol() ==
+               other->result_accuracy().tolerance().rtol() &&
+           result_accuracy().tolerance().ulps() ==
+               other->result_accuracy().tolerance().ulps() &&
+           result_accuracy().mode() == other->result_accuracy().mode();
+  }
+
   void add_single_statistic(Statistic statistic) {
     *mutable_rare()->statistics_viz.add_statistics() = std::move(statistic);
   }

tensorflow/include/tensorflow/compiler/xla/mlir_hlo/_virtual_includes/stablehlo_extension_pass_inc_gen/stablehlo_ext/transforms/passes.h.inc

@@ -3,6 +3,7 @@
 #ifdef GEN_PASS_DECL
 // Generate declarations for all passes.
 #define GEN_PASS_DECL_CHLORECOMPOSEOPSPASS
+#define GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
 #define GEN_PASS_DECL_STABLEHLOCANONICALIZEDYNAMISMPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENENTRYFUNCTIONTUPLESPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENTUPLEPASS
@@ -87,6 +88,82 @@ std::unique_ptr<::mlir::Pass> createChloRecomposeOpsPass() {
 #undef GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 #endif // GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass
+//===----------------------------------------------------------------------===//
+#ifdef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+#undef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#ifdef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
+namespace impl {
+  std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+} // namespace impl
+namespace impl {
+
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Return the dialect that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+private:
+
+  friend std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+    return std::make_unique<DerivedT>();
+  }
+};
+} // namespace impl
+
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+  return impl::createStablehloAddQDQAfterConvPass();
+}
+#undef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass
 //===----------------------------------------------------------------------===//
@@ -360,9 +437,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 
@@ -576,6 +653,23 @@ inline void registerChloRecomposeOpsPassPass() {
   });
 }
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass Registration
+//===----------------------------------------------------------------------===//
+
+inline void registerStablehloAddQDQAfterConvPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
+// Old registration code, kept for temporary backwards compatibility.
+inline void registerStablehloAddQDQAfterConvPassPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass Registration
 //===----------------------------------------------------------------------===//
@@ -684,6 +778,7 @@ inline void registerStablehloRefineShapesPassPass() {
 
 inline void registerPasses() {
   registerChloRecomposeOpsPass();
+  registerStablehloAddQDQAfterConvPass();
   registerStablehloCanonicalizeDynamismPass();
   registerStablehloFlattenEntryFunctionTuplesPass();
   registerStablehloFlattenTuplePass();
@@ -745,6 +840,56 @@ public:
 protected:
 };
 
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Register the dialects that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+};
+
 template <typename DerivedT>
 class StablehloCanonicalizeDynamismPassBase : public ::mlir::OperationPass<func::FuncOp> {
 public:
@@ -907,9 +1052,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 

tensorflow/include/tensorflow/compiler/xla/mlir_hlo/stablehlo_ext/transforms/passes.h.inc

@@ -3,6 +3,7 @@
 #ifdef GEN_PASS_DECL
 // Generate declarations for all passes.
 #define GEN_PASS_DECL_CHLORECOMPOSEOPSPASS
+#define GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
 #define GEN_PASS_DECL_STABLEHLOCANONICALIZEDYNAMISMPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENENTRYFUNCTIONTUPLESPASS
 #define GEN_PASS_DECL_STABLEHLOFLATTENTUPLEPASS
@@ -87,6 +88,82 @@ std::unique_ptr<::mlir::Pass> createChloRecomposeOpsPass() {
 #undef GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 #endif // GEN_PASS_DEF_CHLORECOMPOSEOPSPASS
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass
+//===----------------------------------------------------------------------===//
+#ifdef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+#undef GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DECL_STABLEHLOADDQDQAFTERCONVPASS
+#ifdef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
+namespace impl {
+  std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass();
+} // namespace impl
+namespace impl {
+
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Return the dialect that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+private:
+
+  friend std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+    return std::make_unique<DerivedT>();
+  }
+};
+} // namespace impl
+
+std::unique_ptr<::mlir::Pass> createStablehloAddQDQAfterConvPass() {
+  return impl::createStablehloAddQDQAfterConvPass();
+}
+#undef GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+#endif // GEN_PASS_DEF_STABLEHLOADDQDQAFTERCONVPASS
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass
 //===----------------------------------------------------------------------===//
@@ -360,9 +437,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 
@@ -576,6 +653,23 @@ inline void registerChloRecomposeOpsPassPass() {
   });
 }
 
+//===----------------------------------------------------------------------===//
+// StablehloAddQDQAfterConvPass Registration
+//===----------------------------------------------------------------------===//
+
+inline void registerStablehloAddQDQAfterConvPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
+// Old registration code, kept for temporary backwards compatibility.
+inline void registerStablehloAddQDQAfterConvPassPass() {
+  ::mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return createStablehloAddQDQAfterConvPass();
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // StablehloCanonicalizeDynamismPass Registration
 //===----------------------------------------------------------------------===//
@@ -684,6 +778,7 @@ inline void registerStablehloRefineShapesPassPass() {
 
 inline void registerPasses() {
   registerChloRecomposeOpsPass();
+  registerStablehloAddQDQAfterConvPass();
   registerStablehloCanonicalizeDynamismPass();
   registerStablehloFlattenEntryFunctionTuplesPass();
   registerStablehloFlattenTuplePass();
@@ -745,6 +840,56 @@ public:
 protected:
 };
 
+template <typename DerivedT>
+class StablehloAddQDQAfterConvPassBase : public ::mlir::OperationPass<ModuleOp> {
+public:
+  using Base = StablehloAddQDQAfterConvPassBase;
+
+  StablehloAddQDQAfterConvPassBase() : ::mlir::OperationPass<ModuleOp>(::mlir::TypeID::get<DerivedT>()) {}
+  StablehloAddQDQAfterConvPassBase(const StablehloAddQDQAfterConvPassBase &other) : ::mlir::OperationPass<ModuleOp>(other) {}
+  StablehloAddQDQAfterConvPassBase& operator=(const StablehloAddQDQAfterConvPassBase &) = delete;
+  StablehloAddQDQAfterConvPassBase(StablehloAddQDQAfterConvPassBase &&) = delete;
+  StablehloAddQDQAfterConvPassBase& operator=(StablehloAddQDQAfterConvPassBase &&) = delete;
+  ~StablehloAddQDQAfterConvPassBase() = default;
+
+  /// Returns the command-line argument attached to this pass.
+  static constexpr ::llvm::StringLiteral getArgumentName() {
+    return ::llvm::StringLiteral("stablehlo-ext-add-qdq-after-conv");
+  }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-add-qdq-after-conv"; }
+
+  ::llvm::StringRef getDescription() const override { return "Add quant and dequant ops after convolution op."; }
+
+  /// Returns the derived pass name.
+  static constexpr ::llvm::StringLiteral getPassName() {
+    return ::llvm::StringLiteral("StablehloAddQDQAfterConvPass");
+  }
+  ::llvm::StringRef getName() const override { return "StablehloAddQDQAfterConvPass"; }
+
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const ::mlir::Pass *pass) {
+    return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
+  }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<::mlir::Pass> clonePass() const override {
+    return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
+  }
+
+  /// Register the dialects that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::quant::QuantDialect>();
+    registry.insert<stablehlo::StablehloDialect>();
+  }
+
+  /// Explicitly declare the TypeID for this class. We declare an explicit private
+  /// instantiation because Pass classes should only be visible by the current
+  /// library.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StablehloAddQDQAfterConvPassBase<DerivedT>)
+
+protected:
+};
+
 template <typename DerivedT>
 class StablehloCanonicalizeDynamismPassBase : public ::mlir::OperationPass<func::FuncOp> {
 public:
@@ -907,9 +1052,9 @@ public:
 
   /// Returns the command-line argument attached to this pass.
   static constexpr ::llvm::StringLiteral getArgumentName() {
-    return ::llvm::StringLiteral("legalize-quant-composite");
+    return ::llvm::StringLiteral("stablehlo-ext-legalize-quant-composite");
   }
-  ::llvm::StringRef getArgument() const override { return "legalize-quant-composite"; }
+  ::llvm::StringRef getArgument() const override { return "stablehlo-ext-legalize-quant-composite"; }
 
   ::llvm::StringRef getDescription() const override { return "Lowers the quantization related composites op to native quantized ops."; }
 

tensorflow/include/tensorflow/compiler/xla/pjrt/distributed/client.h

@@ -145,6 +145,11 @@ class DistributedRuntimeClient {
       std::string barrier_id, absl::Duration timeout,
       std::optional<absl::Span<const int32_t>> nodes) = 0;
 
+  // Returns the subset of live nodes. See CoordinationService.GetAliveTasks for
+  // detailed semantics.
+  virtual absl::StatusOr<std::vector<int32_t>> GetLiveNodes(
+      absl::Span<const int32_t> nodes) = 0;
+
   // Returns pointer to coordination service agent, or InternalError if the
   // client does not use coordination service.
   virtual absl::StatusOr<tsl::CoordinationServiceAgent*>

tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client.h

@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/pjrt/gpu/gpu_topology.h"
 #include "xla/pjrt/gpu/gpu_topology.pb.h"
+#include "xla/pjrt/gpu/se_gpu_topology_description.h"
 #include "xla/pjrt/local_device_state.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
@@ -57,98 +58,6 @@ using DeviceTopologyPair =
     std::pair<std::vector<std::unique_ptr<PjRtStreamExecutorDevice>>,
               GpuTopologyProto>;
 
-class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription {
- public:
-  StreamExecutorGpuTopologyDescription(
-      const PjRtPlatformId platform_id, const absl::string_view platform_name,
-      std::shared_ptr<const GpuTopology> gpu_topology,
-      const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& attributes =
-          {},
-      std::optional<stream_executor::GpuTargetConfigProto> target_config =
-          std::nullopt)
-      : platform_id_(platform_id),
-        platform_name_(platform_name),
-        gpu_topology_(std::move(gpu_topology)),
-        attributes_(attributes),
-        target_config_(std::move(target_config)) {}
-
-  bool operator==(const StreamExecutorGpuTopologyDescription& other) const {
-    return this->platform_id() == other.platform_id() &&
-           this->platform_name() == other.platform_name() &&
-           this->platform_version() == other.platform_version() &&
-           this->gpu_topology() == other.gpu_topology();
-  }
-
-  PjRtPlatformId platform_id() const override { return platform_id_; }
-
-  absl::string_view platform_name() const override { return platform_name_; }
-
-  absl::string_view platform_version() const override {
-    return gpu_topology_->platform_version();
-  }
-
-  std::vector<std::unique_ptr<const PjRtDeviceDescription>> DeviceDescriptions()
-      const override {
-    std::vector<std::unique_ptr<const PjRtDeviceDescription>> devices;
-    devices.reserve(gpu_topology_->number_of_devices());
-    for (const int device_id : gpu_topology_->device_ids()) {
-      devices.push_back(std::make_unique<PjRtStreamExecutorDeviceDescription>(
-          device_id, std::string(platform_version())));
-    }
-    return devices;
-  }
-
-  const GpuTopology& gpu_topology() const { return *gpu_topology_; }
-  const GpuTopology* gpu_topology_ptr() const { return gpu_topology_.get(); }
-
-  // No subslice is supported.
-  bool is_subslice_topology() const override { return false; }
-
-  absl::StatusOr<int> ProcessCount() const override {
-    return gpu_topology_->number_of_hosts();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultType() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> LogicalDeviceCountOfDefaultType() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultTypePerProcess() const override {
-    return gpu_topology_->number_of_devices();
-  }
-
-  absl::StatusOr<int> CoreCountOfDefaultTypePerChip() const override {
-    return 1;
-  }
-
-  absl::StatusOr<std::string> Serialize() const override;
-
-  const std::optional<stream_executor::GpuTargetConfigProto>& target_config()
-      const {
-    return target_config_;
-  }
-
-  // Returns vendor specific attributes about the topology.
-  const absl::flat_hash_map<std::string, PjRtDeviceAttribute>& Attributes()
-      const override {
-    return attributes_;
-  }
-
-  absl::StatusOr<Layout> GetDefaultLayout(
-      PrimitiveType element_type,
-      absl::Span<const int64_t> dims) const override;
-
- private:
-  const PjRtPlatformId platform_id_;
-  const std::string platform_name_;
-  std::shared_ptr<const GpuTopology> gpu_topology_;
-  absl::flat_hash_map<std::string, xla::PjRtDeviceAttribute> attributes_;
-  std::optional<stream_executor::GpuTargetConfigProto> target_config_;
-};
-
 class StreamExecutorGpuDevice : public PjRtStreamExecutorDevice {
  public:
   StreamExecutorGpuDevice(int id,