From 2039b482dfe1ea114247e0424675a50162df673b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 1 Sep 2026 15:36:30 -0700 Subject: [PATCH 1/8] Fix alias 'actual' attribute resolving to None in _pywrap_tensorflow for tensorflow_framework. PiperOrigin-RevId: 974744894 --- tensorflow/python/BUILD | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 1ca50ef4d5da69..24e048c55104c1 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -1653,8 +1653,7 @@ pywrap_library( "//conditions:default": "//tensorflow:tf_version_script.lds", }), "tensorflow/tensorflow_framework": select({ - "//tensorflow:windows": None, - "//tensorflow:macos": None, + "//tensorflow:macos": "//tensorflow:tf_exported_symbols.lds", "//conditions:default": "//tensorflow:tf_framework_version_script.lds", }), }, From 73f1a44dd39c22c87c0f2ffff6e17b4b7ba43e8c Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Tue, 1 Sep 2026 16:22:22 -0700 Subject: [PATCH 2/8] [XLA:MSA] Support cross-program prefetch analysis for async ops with cyclic in-place aliasing. PiperOrigin-RevId: 974767010 --- .../memory_space_assignment/algorithm.cc | 36 +++++++++---- .../memory_space_assignment_test.cc | 51 +++++++++++++++++++ 2 files changed, 76 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index 08f5e7630a0b8d..a30bbe7990b0ae 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -322,17 +322,25 @@ bool LooksLikeAnActivation(const HloInstruction* inst, bool permissive_mode) { // Returns true if the use value does not live out of the module. The value // lives out if it is the root or it aliases with another value that lives out. // We recurse to detect the latter case. -bool UseDoesNotLiveOut(const HloUse& use, - const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, - const HloInstruction* root_instruction) { +bool UseDoesNotLiveOut( + const HloUse& use, const HloAliasAnalysis& alias_analysis, + const AliasInfo* alias_info, const HloInstruction* root_instruction, + absl::flat_hash_map& use_to_does_not_live_out) { if (use.instruction == root_instruction && (use.instruction->opcode() == HloOpcode::kTuple || use.instruction->opcode() == HloOpcode::kBitcast)) { return false; } + // If already evaluated or on the current recursion stack, return the result. + auto it = use_to_does_not_live_out.find(use); + if (it != use_to_does_not_live_out.end()) { + return it->second; + } + // Mark this use as in-progress (defaulting to true) to break in-place + // aliasing cycles (e.g. async start <-> done). + use_to_does_not_live_out[use] = true; auto in_place_pairs = alias_info->GetInPlaceInputOutputPairs(use.instruction); - return absl::c_all_of( + bool does_not_live_out = absl::c_all_of( in_place_pairs, [&](const std::pair& in_place_pair) { if (in_place_pair.first.operand_number == use.operand_number && @@ -346,13 +354,17 @@ bool UseDoesNotLiveOut(const HloUse& use, .GetUses()) { if (nested_use != use && !UseDoesNotLiveOut(nested_use, alias_analysis, alias_info, - root_instruction)) { + root_instruction, + use_to_does_not_live_out)) { return false; } } } return true; }); + // Update the map with the finalized result. + use_to_does_not_live_out[use] = does_not_live_out; + return does_not_live_out; } // Filters out buffer uses that cannot use the cross-program prefetch due to @@ -370,11 +382,13 @@ std::vector FindCrossProgramPrefetchUses( ->entry_computation() ->root_instruction(); - absl::c_copy_if(buffer_uses, std::back_inserter(uses), - [&](const HloUse& use) { - return UseDoesNotLiveOut(use, alias_analysis, alias_info, - root_instruction); - }); + absl::flat_hash_map use_to_does_not_live_out; + for (const HloUse& use : buffer_uses) { + if (UseDoesNotLiveOut(use, alias_analysis, alias_info, root_instruction, + use_to_does_not_live_out)) { + uses.push_back(use); + } + } return uses; } diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc index bfbc7d20331b0a..7500e842ea22eb 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc @@ -19562,6 +19562,57 @@ ENTRY main { }, preset_assignments.get(), alias_analysis.get())); } + +TEST_F(MemorySpaceAssignmentTest, AsyncBarrierAliasingCycle) { + absl::string_view hlo_string = R"hlo( +HloModule sc_async_barrier_aliasing, is_scheduled=true + + sc_custom_call_comp { + ROOT %custom_call = (f32[32,128,8]{1,2,0}, u32[], u32[]) custom-call(), + custom_call_target="SparseCoreBarrierStart", + custom_call_has_side_effect=true + } + + sc_main { + %input = f32[32,128,8]{1,2,0} parameter(0) + %output = f32[32,128,8]{1,2,0} parameter(1) + %send_sflag = u32[] parameter(2) + %recv_sflag = u32[] parameter(3) + ROOT %a2a-result = f32[32,128,8]{1,2,0} custom-call(%input, %output, %send_sflag, %recv_sflag), + custom_call_target="SparseCoreBarrierDoneAndCollective", + output_to_operand_aliasing={{}: (1, {})} + } + +ENTRY %Comp_spmd { + %p0 = f32[32,128,8]{1,2,0} parameter(0) + %p1 = f32[32,128,8]{1,2,0} parameter(1) + + %sc_custom_call_start = ((), (f32[32,128,8]{1,2,0}, u32[], u32[]), s32[]) call-start(), + to_apply=%sc_custom_call_comp + %sc_custom_call_done = (f32[32,128,8]{1,2,0}, u32[], u32[]) call-done(%sc_custom_call_start) + %send_sflag = u32[] get-tuple-element(%sc_custom_call_done), index=1 + %recv_sflag = u32[] get-tuple-element(%sc_custom_call_done), index=2 + + %sc_start = ((f32[32,128,8]{1,2,0}, f32[32,128,8]{1,2,0}, u32[], u32[]), f32[32,128,8]{1,2,0}, s32[]) + call-start(%p0, %p1, %send_sflag, %recv_sflag), + to_apply=%sc_main, + output_to_operand_aliasing={ + {0,1}: (1, {}), // update output buffer at {0,0} + {0,2}: (2, {}), // update send_sflag + {0,3}: (3, {}), // update recv_sflag + {1}: (1, {})} // update output buffer at {1} + %sc_output = f32[32,128,8]{1,2,0} call-done(%sc_start) + %neg = f32[32,128,8]{1,2,0} negate(%sc_output) + + ROOT %root = (f32[32,128,8]{1,2,0}, f32[32,128,8]{1,2,0}) tuple(%neg, %p0) +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_string)); + EXPECT_NO_FATAL_FAILURE(AssignMemorySpace(module.get())); +} + } // namespace } // namespace memory_space_assignment From 0b08745154d434a9e5f247fb25bfba70af71f19b Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Tue, 1 Sep 2026 16:32:02 -0700 Subject: [PATCH 3/8] Add TFLiteConverter.from_mlir_bytecode in tf.lite Python API. PiperOrigin-RevId: 974771869 --- tensorflow/compiler/mlir/lite/python/BUILD | 1 - .../mlir/lite/python/slim_model_importer.cc | 1 + tensorflow/lite/python/BUILD | 2 +- tensorflow/lite/python/convert.py | 47 +++++++++++++++++ tensorflow/lite/python/lite.py | 52 ++++++++++++++++++- 5 files changed, 100 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 8411fbe7b0e56a..5ed120ebbeade1 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -237,7 +237,6 @@ cc_library( "@llvm-project//mlir:IR", "@llvm-project//mlir:Parser", "@llvm-project//mlir:Pass", - "@llvm-project//mlir:ReconcileUnrealizedCasts", "@llvm-project//mlir:Support", "@stablehlo//:stablehlo_ops", "@stablehlo//:stablehlo_passes", diff --git a/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc b/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc index b8712b51d11887..998ad8f15f92b3 100644 --- a/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc +++ b/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc @@ -222,6 +222,7 @@ absl::StatusOr> LoadSlimModel( mlir::PassManager pm(context); pm.addPass(mlir::odml::CreateDropShapeAssertionsPass()); + pm.addPass(mlir::odml::CreateLegalizeVhloQuantCustomCallsPass()); pm.addPass(mlir::stablehlo::createVhloLegalizeToStablehloPass()); if (mlir::failed(pm.run(*module))) { return absl::InternalError("Failed to legalize VHLO to StableHLO."); diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD index b095448a5f7037..a9cf236970aac8 100644 --- a/tensorflow/lite/python/BUILD +++ b/tensorflow/lite/python/BUILD @@ -227,7 +227,7 @@ py_library( ":util", "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py", "//tensorflow/compiler/mlir/quantization/tensorflow/python:representative_dataset", - "//tensorflow/core:protos_all_py", + "//tensorflow/core/framework:graph_proto_py_proto", "//tensorflow/lite/experimental/microfrontend:audio_microfrontend_py", "//tensorflow/lite/profiling/proto:model_runtime_info_py", "//tensorflow/lite/profiling/proto:profiling_info_py", diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 42515e633f9193..9d319e54b894b4 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -348,6 +348,38 @@ def convert( raise converter_error +def convert_mlir_bytecode( + conversion_flags: _conversion_flags_pb2.ConverterFlags, + model_dir: str, + output_file_path: str, +): + """Converts `model_dir` to a TFLite model file directly. + + Args: + conversion_flags: Proto describing conversion properties, see + `compiler/mlir/lite/converter_flags.proto`. + model_dir: Directory containing the MLIR bytecode and weights. + output_file_path: Path where the TFLite model should be saved. + + Returns: + Status or result of the conversion. + + Raises: + ConverterError: When conversion fails. + """ + try: + return wrap_converter.wrapped_convert_mlir_bytecode( + conversion_flags.SerializeToString(), + model_dir, + output_file_path, + ) + except Exception as e: + converter_error = ConverterError(str(e)) + for error_data in _metrics_wrapper.retrieve_collected_errors(): + converter_error.append_error(error_data) + raise converter_error from e + + def build_model_flags( change_concat_input_ranges=False, allow_nonexistent_arrays=False, @@ -419,6 +451,7 @@ def build_conversion_flags( accumulation_type=None, allow_bfloat16=False, unfold_large_splat_constant=False, + fold_fp16_resource_casts=True, supported_backends=None, disable_per_channel_quantization=False, enable_mlir_dynamic_range_quantizer=False, @@ -449,6 +482,8 @@ def build_conversion_flags( serialize_debug_metadata=False, unsafe_fuse_dynamic_shaped_broadcast=False, unsafe_single_batch_rank_reduction=False, + enable_debug=False, + debug_dir=None, **_, ): """Builds protocol buffer describing a conversion of a model. @@ -519,6 +554,8 @@ def build_conversion_flags( inference with the bfloat16 type. unfold_large_splat_constant: Whether to unfold large splat constant tensors in the flatbuffer model to reduce size. + fold_fp16_resource_casts: Whether to fold 16-bit float (fp16/bf16) resource + casts. supported_backends: List of TFLite backends which needs to check compatibility. disable_per_channel_quantization: Disable per-channel quantized weights for @@ -589,6 +626,9 @@ def build_conversion_flags( the source model. unsafe_single_batch_rank_reduction: When set to true, enable the unsafe single batch rank reduction. + enable_debug: When set to true, enable debug mode. + debug_dir: Directory to save debug output. + **_: Additional unused keyword arguments. Returns: conversion_flags: protocol buffer describing the conversion process. @@ -653,6 +693,8 @@ def build_conversion_flags( ) conversion_flags.allow_bfloat16 = allow_bfloat16 conversion_flags.unfold_large_splat_constant = unfold_large_splat_constant + if hasattr(conversion_flags, "fold_fp16_resource_casts"): + conversion_flags.fold_fp16_resource_casts = fold_fp16_resource_casts if supported_backends: conversion_flags.supported_backends.extend(supported_backends) conversion_flags.disable_per_channel_quantization = ( @@ -696,6 +738,11 @@ def build_conversion_flags( elide_elementsattrs_if_larger ) + if hasattr(conversion_flags, "enable_debug"): + conversion_flags.enable_debug = enable_debug + if debug_dir is not None and hasattr(conversion_flags, "debug_dir"): + conversion_flags.debug_dir = debug_dir + if use_buffer_offset is not None: conversion_flags.use_buffer_offset = use_buffer_offset if reduce_type_precision is not None: diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index bd9af24daa0491..0c6dfa0c9adbb9 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -16,6 +16,7 @@ import enum import functools +import os import pprint import shutil import sys @@ -41,6 +42,7 @@ from tensorflow.lite.python.convert import convert_graphdef as _convert_graphdef from tensorflow.lite.python.convert import convert_graphdef_with_arrays as _convert_graphdef_with_arrays from tensorflow.lite.python.convert import convert_jax_hlo as _convert_jax_hlo +from tensorflow.lite.python.convert import convert_mlir_bytecode as _convert_mlir_bytecode from tensorflow.lite.python.convert import convert_saved_model as _convert_saved_model from tensorflow.lite.python.convert import ConverterError # pylint: disable=unused-import from tensorflow.lite.python.convert import deduplicate_readonly_buffers as _deduplicate_readonly_buffers @@ -641,6 +643,7 @@ def __init__(self): self._experimental_lower_tensor_list_ops = True self._experimental_default_to_single_batch_in_tensor_list_ops = False self._experimental_unfold_large_splat_constant = False + self._experimental_fold_fp16_resource_casts = True self._experimental_tf_quantization_mode = None # If unset, bias:int32 is by default except 16x8 quant. # For 16x8 quant, bias:int64 is used to prevent any overflow by default. @@ -697,6 +700,8 @@ def __init__(self): self.print_ir_module_scope = None self.elide_elementsattrs_if_larger = None self.serialize_debug_metadata = False + self.enable_debug = False + self.debug_dir = None def _grappler_config(self, optimizers=None): """Creates a tf.compat.v1.ConfigProto for configuring Grappler. @@ -818,6 +823,7 @@ def _get_base_converter_args(self): "unfold_large_splat_constant": ( self._experimental_unfold_large_splat_constant ), + "fold_fp16_resource_casts": self._experimental_fold_fp16_resource_casts, "default_to_single_batch_in_tensor_list_ops": ( self._experimental_default_to_single_batch_in_tensor_list_ops ), @@ -843,6 +849,8 @@ def _get_base_converter_args(self): "print_ir_after": self.print_ir_after, "print_ir_module_scope": self.print_ir_module_scope, "elide_elementsattrs_if_larger": self.elide_elementsattrs_if_larger, + "enable_debug": self.enable_debug, + "debug_dir": self.debug_dir, "use_buffer_offset": self._experimental_use_buffer_offset, "reduce_type_precision": self._experimental_reduce_type_precision, "use_stablehlo_quantizer": self.experimental_use_stablehlo_quantizer, @@ -1219,7 +1227,12 @@ def _convert_and_export_metrics(self, convert_func, *args, **kwargs): self._increase_conversion_success_metric() self._set_conversion_latency_metric(round(elapsed_time_ms)) self._tflite_metrics.export_metrics() - if self.exclude_conversion_metadata or self._experimental_use_buffer_offset: + if ( + self.exclude_conversion_metadata + or self._experimental_use_buffer_offset + or result is None + or isinstance(result, (str, os.PathLike)) + ): return result # TODO(b/286886803): add support for adding user metadata with # use_buffer_offset flags @@ -2134,6 +2147,30 @@ def convert(self): ) +class TFLiteMlirBytecodeConverterV2(TFLiteConverterBaseV2): + """Converts the given MLIR bytecode into TensorFlow Lite model.""" + + def __init__(self, model_dir): + """Constructor for TFLiteConverter. + + Args: + model_dir: Directory containing the MLIR bytecode and weights. + """ + super(TFLiteMlirBytecodeConverterV2, self).__init__() + self._model_dir = model_dir + + @_export_metrics + def convert(self, output_file_path=None): + """Converts the MLIR bytecode and saves directly to output_file_path without heap RAM spike.""" + converter_kwargs = self._get_base_converter_args() + # Build conversion flags. + conversion_flags = _build_conversion_flags(**converter_kwargs) + if output_file_path is None: + output_file_path = os.path.join(self._model_dir, "model.tflite") + _convert_mlir_bytecode(conversion_flags, self._model_dir, output_file_path) + return output_file_path + + @_tf_export("lite.TFLiteConverter", v1=[]) class TFLiteConverterV2(TFLiteFrozenGraphConverterV2): """Converts a TensorFlow model into TensorFlow Lite model. @@ -2390,6 +2427,19 @@ def experimental_from_jax(cls, serving_funcs, inputs): # pylint: enable=protected-access return TFLiteJaxConverterV2(serving_funcs, inputs) + @classmethod + def _from_mlir_bytecode(cls, model_dir): + """Creates a TFLiteConverter object from MLIR bytecode. + + Args: + model_dir: Directory containing the MLIR bytecode and weights. + + Returns: + TFLiteConverter object. + """ + + return TFLiteMlirBytecodeConverterV2(model_dir) + # pylint: disable=useless-super-delegation def convert(self): """Converts a TensorFlow GraphDef based on instance variables. From 08d66ff30ce61b09ac5ab0fffec08b1ab088ba65 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 1 Sep 2026 17:37:06 -0700 Subject: [PATCH 4/8] Introduce PositionalWeightedCombiner for SparseCore embeddings. PiperOrigin-RevId: 974800503 --- .../python/tpu/tpu_embedding_v2_utils.py | 74 +++++++++++- .../python/tpu/tpu_embedding_v2_utils_test.py | 105 ++++++++++++++++++ ...edding.-positional-weighted-combiner.pbtxt | 18 +++ ...ensorflow.tpu.experimental.embedding.pbtxt | 4 + ...edding.-positional-weighted-combiner.pbtxt | 18 +++ ...ensorflow.tpu.experimental.embedding.pbtxt | 4 + 6 files changed, 220 insertions(+), 3 deletions(-) create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt diff --git a/tensorflow/python/tpu/tpu_embedding_v2_utils.py b/tensorflow/python/tpu/tpu_embedding_v2_utils.py index a6f76120c216c2..d01ec7cd2345d2 100644 --- a/tensorflow/python/tpu/tpu_embedding_v2_utils.py +++ b/tensorflow/python/tpu/tpu_embedding_v2_utils.py @@ -1251,6 +1251,71 @@ def __str__(self) -> str: return self.combiner +@tf_export("tpu.experimental.embedding.PositionalWeightedCombiner") +class PositionalWeightedCombiner(_WithSlotVariables): + """Positional weighted combiner for TPU embeddings on SparseCore. + + This class specifies a positional weighted sum combiner over the TPU embedding + lookup results. It also supports a custom optimizer computation to be + performed on the positional weights. + """ + + def __init__( + self, + max_valency: int, + initializer: init_ops_v2.Initializer, + custom_optimizer: CustomOptimizer, + ) -> Any: + """Initializes the positional weighted combiner. + + Args: + max_valency: The valency cap of the corresponding feature. This value is + also used as the number of positional weights. + initializer: The initializer for the positional weights. + custom_optimizer: CustomOptimizer encapsulating the custom optimizer + computation, learning rate, slot names, slot initializers, and + hyperparameters. + """ + super().__init__() + self.combiner = "positional_weighted" + self.max_valency = max_valency + self.initializer = initializer + self.custom_optimizer = custom_optimizer + + if max_valency <= 0: + raise ValueError(f"Expect max_valency > 0, but got {max_valency}.") + if initializer is None: + raise ValueError("Expect initializer to be not None.") + if custom_optimizer is None or not isinstance( + custom_optimizer, CustomOptimizer + ): + raise ValueError( + "Expect custom_optimizer to be an instance of CustomOptimizer, but" + f" got {type(custom_optimizer)}." + ) + + self.combiner_weights_learning_rate = custom_optimizer.learning_rate + self.custom_optimizer_function = custom_optimizer.custom_computation + self.custom_computation = custom_optimizer.custom_computation + + self._slot_names_attr = tuple(custom_optimizer._slot_names()) + self._slot_initializers_attr = tuple(custom_optimizer._slot_initializers()) + self._hyperparameters_attr = tuple(custom_optimizer.hyperparameters or ()) + + def _slot_names(self) -> List[Text]: + return list(self._slot_names_attr) + + def _slot_initializers(self) -> List[init_ops_v2.Initializer]: + return list(self._slot_initializers_attr) + + @property + def hyperparameters(self) -> List[Union[float, Callable[[], float]]]: + return list(self._hyperparameters_attr) + + def __str__(self) -> str: + return self.combiner + + @tf_export("tpu.experimental.embedding.QuantizationConfig") class QuantizationConfig: """Settings for simulated quantization of the tpu embedding table. @@ -1432,10 +1497,13 @@ def __init__( f"String argument `combiner` must be in {accepted_str_combiners}. " f"Received: {combiner}") - elif not isinstance(combiner, CustomCombiner): + elif not isinstance( + combiner, (CustomCombiner, PositionalWeightedCombiner) + ): raise ValueError( - f"Argument `combiner` should either be a str or a CustomCombiner. " - f"Received: {type(combiner)}" + "Argument `combiner` should either be a str, CustomCombiner, or" + " PositionalWeightedCombiner." + f" Received: {type(combiner)}" ) if name is None: diff --git a/tensorflow/python/tpu/tpu_embedding_v2_utils_test.py b/tensorflow/python/tpu/tpu_embedding_v2_utils_test.py index e0ba452d68802a..7efb075157b3d9 100644 --- a/tensorflow/python/tpu/tpu_embedding_v2_utils_test.py +++ b/tensorflow/python/tpu/tpu_embedding_v2_utils_test.py @@ -232,6 +232,111 @@ def test_sort_device_spec_strings(self): self.assertEqual(sorted_specs, sorted(device_spec_strings)) +class PositionalWeightedCombinerTest(test.TestCase): + + def test_custom_optimizer_validation(self): + def dummy_fn(grad, var, slots, lr, hyperparams): + return var - grad * lr + + # Valid config + optimizer = tpu_embedding_v2_utils.CustomOptimizer( + custom_computation=dummy_fn, + slot_names=['slot_a'], + slot_initializers=[init_ops_v2.Zeros()], + ) + self.assertEqual(optimizer._slot_names(), ['slot_a']) + + # slot_names is None, slot_initializers is not None + with self.assertRaisesRegex(ValueError, 'must match'): + tpu_embedding_v2_utils.CustomOptimizer( + custom_computation=dummy_fn, + slot_names=None, + slot_initializers=[init_ops_v2.Zeros()], + ) + + # slot_names is not None, slot_initializers is None + with self.assertRaisesRegex(ValueError, 'must match'): + tpu_embedding_v2_utils.CustomOptimizer( + custom_computation=dummy_fn, + slot_names=['slot_a'], + slot_initializers=None, + ) + + # length mismatch + with self.assertRaisesRegex(ValueError, 'must match'): + tpu_embedding_v2_utils.CustomOptimizer( + custom_computation=dummy_fn, + slot_names=['slot_a', 'slot_b'], + slot_initializers=[init_ops_v2.Zeros()], + ) + + def test_positional_weighted_combiner(self): + def dummy_fn(grad, var, slots, lr, hyperparams): + return var - grad * lr + + optimizer = tpu_embedding_v2_utils.CustomOptimizer( + custom_computation=dummy_fn, + learning_rate=0.1, + slot_names=['slot_a'], + slot_initializers=[init_ops_v2.Zeros()], + hyperparameters=[0.5], + ) + combiner = tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=10, + initializer=init_ops_v2.Ones(), + custom_optimizer=optimizer, + ) + self.assertEqual(combiner.max_valency, 10) + self.assertEqual(combiner.custom_optimizer, optimizer) + self.assertEqual(combiner.combiner_weights_learning_rate, 0.1) + self.assertEqual(combiner.custom_optimizer_function, dummy_fn) + self.assertEqual(combiner.custom_computation, dummy_fn) + self.assertEqual(combiner._slot_names(), ['slot_a']) + self.assertLen(combiner._slot_initializers(), 1) + self.assertEqual(combiner.hyperparameters, [0.5]) + self.assertEqual(str(combiner), 'positional_weighted') + + # equality + combiner_same = tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=10, + initializer=combiner.initializer, + custom_optimizer=optimizer, + ) + self.assertEqual(combiner, combiner_same) + + # invalid max_valency <= 0 + with self.assertRaisesRegex(ValueError, 'max_valency > 0'): + tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=0, + initializer=init_ops_v2.Ones(), + custom_optimizer=optimizer, + ) + + # invalid initializer is None + with self.assertRaisesRegex(ValueError, 'initializer to be not None'): + tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=10, + initializer=None, + custom_optimizer=optimizer, + ) + + # custom_optimizer is not an instance of CustomOptimizer + with self.assertRaisesRegex(ValueError, 'CustomOptimizer'): + tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=10, + initializer=init_ops_v2.Ones(), + custom_optimizer='invalid_optimizer', + ) + + # custom_optimizer is None + with self.assertRaisesRegex(ValueError, 'CustomOptimizer'): + tpu_embedding_v2_utils.PositionalWeightedCombiner( + max_valency=10, + initializer=init_ops_v2.Ones(), + custom_optimizer=None, + ) + + if __name__ == '__main__': v2_compat.enable_v2_behavior() test.main() diff --git a/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt new file mode 100644 index 00000000000000..dbb47eb9ef94d5 --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt @@ -0,0 +1,18 @@ +path: "tensorflow.tpu.experimental.embedding.PositionalWeightedCombiner" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "hyperparameters" + mtype: "" + } + member_method { + name: "__eq__" + argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "__init__" + argspec: "args=[\'self\', \'max_valency\', \'initializer\', \'custom_optimizer\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.pbtxt index 0d5008aaf46e88..b7ec22419d5e4e 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.tpu.experimental.embedding.pbtxt @@ -28,6 +28,10 @@ tf_module { name: "FeatureConfig" mtype: "" } + member { + name: "PositionalWeightedCombiner" + mtype: "" + } member { name: "QuantizationConfig" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt new file mode 100644 index 00000000000000..dbb47eb9ef94d5 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.-positional-weighted-combiner.pbtxt @@ -0,0 +1,18 @@ +path: "tensorflow.tpu.experimental.embedding.PositionalWeightedCombiner" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "hyperparameters" + mtype: "" + } + member_method { + name: "__eq__" + argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "__init__" + argspec: "args=[\'self\', \'max_valency\', \'initializer\', \'custom_optimizer\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.pbtxt index 0d5008aaf46e88..b7ec22419d5e4e 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.tpu.experimental.embedding.pbtxt @@ -28,6 +28,10 @@ tf_module { name: "FeatureConfig" mtype: "" } + member { + name: "PositionalWeightedCombiner" + mtype: "" + } member { name: "QuantizationConfig" mtype: "" From 36072d0bd8c26388fe55029e37cf63740a64c906 Mon Sep 17 00:00:00 2001 From: Zac Mustin Date: Tue, 1 Sep 2026 17:51:50 -0700 Subject: [PATCH 5/8] Rollback of PR #47176 Causes internal build breakage. Reverts 8eee66319446737cdb639979c88a4a62903eba63 PiperOrigin-RevId: 974806297 --- .../xla/xla/backends/gpu/collectives/BUILD | 36 +-- .../gpu/collectives/mori_communicator.cc | 286 ++++++------------ .../gpu/collectives/mori_communicator.h | 79 +++-- .../gpu/collectives/mori_kernels.cu.cc | 19 -- .../backends/gpu/collectives/mori_kernels.h | 27 -- .../xla/backends/gpu/collectives/mori_stub.h | 75 ----- 6 files changed, 132 insertions(+), 390 deletions(-) delete mode 100644 third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc delete mode 100644 third_party/xla/xla/backends/gpu/collectives/mori_kernels.h diff --git a/third_party/xla/xla/backends/gpu/collectives/BUILD b/third_party/xla/xla/backends/gpu/collectives/BUILD index b21fc137166589..6e6aadc1c6ff5e 100644 --- a/third_party/xla/xla/backends/gpu/collectives/BUILD +++ b/third_party/xla/xla/backends/gpu/collectives/BUILD @@ -1,8 +1,4 @@ -load( - "@local_config_rocm//rocm:build_defs.bzl", - "if_rocm_is_configured", - "rocm_library", -) +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl_is_configured") load("//xla:xla.default.bzl", "xla_cc_test") load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm_is_configured") @@ -1271,26 +1267,6 @@ cc_library( ], ) -rocm_library( - name = "mori_kernels", - srcs = [ - "mori_kernels.cu.cc", - ], - hdrs = [ - "mori_kernels.h", - "mori_stub.h", - ], - # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], - copts = ["-U__HIP_DISABLE_CPP_FUNCTIONS__"], # <-- only if needed - linkstatic = True, - tags = [ - "gpu", - "no-oneapi", - "rocm-only", - ], - deps = [], -) - cc_library( name = "mori_collectives", srcs = [ @@ -1300,13 +1276,10 @@ cc_library( hdrs = [ "mori_collectives.h", "mori_communicator.h", - "mori_kernels.h", "mori_stub.h", ], tags = [ "gpu", - "no-oneapi", - "rocm-only", ], visibility = ["//visibility:public"], deps = [ @@ -1326,7 +1299,6 @@ cc_library( "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", "//xla/core/collectives:reduction_kind", - "//xla/core/collectives:symmetric_memory", "//xla/pjrt/distributed:key_value_store_interface", "//xla/runtime:device_id", "//xla/runtime:process_id", @@ -1335,7 +1307,6 @@ cc_library( "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream", "//xla/stream_executor:stream_executor_h", - "//xla/stream_executor/rocm:rocm_status", "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", @@ -1358,10 +1329,7 @@ cc_library( "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", "@tsl//tsl/platform:numbers", - ] + if_rocm_is_configured([ - ":mori_kernels", - "@local_config_rocm//rocm:rocm_headers", - ]), + ], alwayslink = True, ) diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc index 8dfb3bfa4e80e9..d1d71dc7fd55c8 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc +++ b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc @@ -36,14 +36,13 @@ limitations under the License. #include "xla/backends/gpu/collectives/cancellation_token.h" #include "xla/backends/gpu/collectives/gpu_collectives.h" #include "xla/backends/gpu/collectives/mori_collectives.h" -#include "xla/backends/gpu/collectives/mori_kernels.h" +#include "xla/backends/gpu/collectives/mori_stub.h" #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/reduction_kind.h" #include "xla/future.h" #include "xla/primitive_util.h" #include "xla/stream_executor/device_address.h" -#include "xla/stream_executor/rocm/rocm_status.h" #include "xla/stream_executor/stream.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -52,104 +51,31 @@ limitations under the License. namespace shmem = ::mori::shmem; namespace xla::gpu { -using ::mori::collective::CollectivesFacade; -namespace { - -hipStream_t AsHipStream(se::Stream* stream) { - return reinterpret_cast( +static auto AsRocmStream(se::Stream* stream) { + return reinterpret_cast( stream->platform_specific_handle().stream); } -size_t ToMoriByteCount(PrimitiveType dtype, size_t count) { +static size_t ToMoriByteCount(PrimitiveType dtype, size_t count) { if (primitive_util::IsComplexType(dtype)) { count *= 2; } return count * primitive_util::BitWidth(dtype) / 8; } -absl::StatusOr<::mori::collective::DataType> ToMoriDataType( - PrimitiveType dtype) { -#define MORI_TYPE_DISPATCH(x) \ - case x: \ - return ::mori::collective::DataType::x; - switch (dtype) { - MORI_TYPE_DISPATCH(F8E5M2) - MORI_TYPE_DISPATCH(F8E4M3FN) - MORI_TYPE_DISPATCH(F16) - MORI_TYPE_DISPATCH(BF16) - MORI_TYPE_DISPATCH(S8) - MORI_TYPE_DISPATCH(U8) - MORI_TYPE_DISPATCH(S32) - MORI_TYPE_DISPATCH(U32) - MORI_TYPE_DISPATCH(S64) - MORI_TYPE_DISPATCH(U64) - MORI_TYPE_DISPATCH(F32) - MORI_TYPE_DISPATCH(F64) - default: - return absl::UnimplementedError(absl::StrFormat( - "MORI: unsupported dtype: %d", static_cast(dtype))); - } -#undef MORI_TYPE_DISPATCH -} - -// Translate an XLA ReductionKind to the facade's reduction-op enum. -absl::StatusOr<::mori::collective::ReduceOpKind> ToMoriReduceOp( - ReductionKind r) { -#define MORI_OP_DISPATCH(x) \ - case ReductionKind::x: \ - return ::mori::collective::ReduceOpKind::x; - switch (r) { - MORI_OP_DISPATCH(SUM) - MORI_OP_DISPATCH(PRODUCT) - MORI_OP_DISPATCH(MIN) - MORI_OP_DISPATCH(MAX) - default: - return absl::UnimplementedError(absl::StrFormat( - "MORI: unsupported reduction op: %d", static_cast(r))); - } -#undef MORI_OP_DISPATCH -} - -absl::StatusOr ToStream(const Communicator::Executor& executor) { - if (auto* gpu_executor = - absl::down_cast(&executor)) { - return gpu_executor->stream(); - } - return InvalidArgument("Communicator executor is not a GPU executor"); -} -} // namespace - absl::StatusOr> MoriCommunicator::Create( MoriCollectives* coll, std::shared_ptr cancel, int rank, absl::Span rank_to_pe) { auto comm = absl::WrapUnique(new MoriCommunicator(coll, cancel)); const int num_ranks = static_cast(rank_to_pe.size()); - if (num_ranks <= 0) { - return absl::InvalidArgumentError(absl::StrFormat( - "MoriCommunicator: unsupported number of ranks %d", num_ranks)); - } comm->rank_ = rank; comm->num_ranks_ = num_ranks; - - // The CollectivesFacade owns this communicator's symmetric-heap staging - // buffer and the push reduce-scatter group counters. It records the rank - // identity (rank/num_ranks) and allocates the ~2GB staging; the unique_ptr - // frees it (before ShmemFinalize) when the communicator is destroyed. - const size_t buffer_size = 2UL << 30; // 2GB - comm->facade_ = CollectivesFacade::Create(rank, num_ranks, buffer_size); - if (comm->facade_ == nullptr) { - return absl::InternalError("CollectivesFacade::Create failed"); - } VLOG(1) << "Created " << *comm << " with participants: " << num_ranks; return comm; } -MoriCommunicator::~MoriCommunicator() { - // facade_ (unique_ptr) releases this communicator's staging + counters via - // the CollectivesFacade dtor here, before MoriCollectives::Finalize() -> - // ShmemFinalize. -} +MoriCommunicator::~MoriCommunicator() {} #define CHECK_CANCELLED() \ if (cancel_->IsCancelled()) { \ @@ -171,26 +97,41 @@ absl::Status MoriCommunicator::Abort() { return absl::OkStatus(); } -absl::Status MoriCommunicator::Barrier(const Executor& executor) { +absl::Status MoriCommunicator::Barrier(const Communicator::Executor& executor) { VLOG(1) << "Barrier: " << ToString(); CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - return se::gpu::ToStatus(facade_->RunBarrier(AsHipStream(stream))); + (void)stream; + // return xla_mori::BarrierOnStream(AsRocmStream(stream)); + return absl::OkStatus(); } absl::StatusOr MoriCommunicator::NumRanks() const { + VLOG(5) << "Get the number of ranks in MORI communicator: " << ToString(); CHECK_CANCELLED() + return static_cast(num_ranks_); } absl::StatusOr MoriCommunicator::CurrentRank() { + VLOG(5) << "Get current rank in MORI communicator: " << ToString(); CHECK_CANCELLED() + return static_cast(rank_); } std::string MoriCommunicator::ToString() const { - return absl::StrFormat("MoriCommunicator(rank=%d, num_ranks=%d)", rank_, - num_ranks_); + return absl::StrFormat("MoriCommunicator(rank=%d, num_ranks=%d, my_pe=%d)", + rank_, num_ranks_, shmem::ShmemMyPe()); +} + +absl::StatusOr MoriCommunicator::ToStream( + const Executor& executor) { + if (auto* gpu_executor = + absl::down_cast(&executor)) { + return gpu_executor->stream(); + } + return InvalidArgument("Communicator executor is not a GPU executor"); } Future<> MoriCommunicator::AllReduce(se::DeviceAddressBase send_buffer, @@ -260,20 +201,20 @@ Future<> MoriCommunicator::CollectivePermute( }); } -Future<> MoriCommunicator::Send(se::DeviceAddressBase send_buffer, +Future<> MoriCommunicator::Send(se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) { - return Execute([send_buffer, dtype, count, peer, &executor, this]() { - return LaunchSend(send_buffer, dtype, count, peer, executor); - }); + return P2P(P2PType::Send, dtype, recv_buffer, send_buffer, count, peer, + executor); } Future<> MoriCommunicator::Recv(se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) { - return Execute([recv_buffer, dtype, count, peer, &executor, this]() { - return LaunchRecv(recv_buffer, dtype, count, peer, executor); - }); + return P2P(P2PType::Recv, dtype, recv_buffer, send_buffer, count, peer, + executor); } absl::Status MoriCommunicator::LaunchAllGather( @@ -281,14 +222,11 @@ absl::Status MoriCommunicator::LaunchAllGather( PrimitiveType dtype, size_t count, const Executor& executor) { CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - - VLOG(3) << "Launch AllGather: send_buffer=" << send_buffer.opaque() + VLOG(3) << "LaunchAllGather: send_buffer=" << send_buffer.opaque() << " recv_buffer=" << recv_buffer.opaque() << " count=" << count << " dtype=" << primitive_util::LowercasePrimitiveTypeName(dtype) - << " stream=" << AsHipStream(stream); - return se::gpu::ToStatus(facade_->RunAllGather( - send_buffer.opaque(), recv_buffer.opaque(), ToMoriByteCount(dtype, count), - AsHipStream(stream))); + << " stream=" << AsRocmStream(stream); + return absl::UnimplementedError("Not implemented"); } absl::Status MoriCommunicator::LaunchAllReduce( @@ -296,20 +234,25 @@ absl::Status MoriCommunicator::LaunchAllReduce( PrimitiveType dtype, size_t count, ReductionKind reduction_kind, const Executor& executor) { CHECK_CANCELLED() + ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + auto gpu_stream = AsRocmStream(stream); + (void)gpu_stream; + void* source_ptr = send_buffer.opaque(); + void* dest_ptr = recv_buffer.opaque(); + (void)source_ptr; + (void)dest_ptr; + if (primitive_util::IsComplexType(dtype)) { + count *= 2; + } VLOG(3) << absl::StreamFormat( - "Launch AllReduce: send_buffer=%p; recv_buffer=%p; dtype=%s; count=%d; " - "reduction_kind=%v; stream=%p", + "Launch MORI AllReduce send_buffer=%p; recv_buffer=%p; dtype=%s; " + "count=%d; reduction_kind=%v; device_ordinal=%d", send_buffer.opaque(), recv_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), count, reduction_kind, - stream); - - ABSL_ASSIGN_OR_RETURN(auto dt, ToMoriDataType(dtype)); - ABSL_ASSIGN_OR_RETURN(auto op, ToMoriReduceOp(reduction_kind)); - return se::gpu::ToStatus(facade_->RunAllReduce(send_buffer.opaque(), - recv_buffer.opaque(), count, - dt, op, AsHipStream(stream))); + stream->parent()->device_ordinal()); + return absl::UnimplementedError("Not implemented"); } absl::Status MoriCommunicator::LaunchReduceScatter( @@ -322,78 +265,8 @@ absl::Status MoriCommunicator::LaunchReduceScatter( VLOG(3) << "LaunchReduceScatter: send_buffer=" << send_buffer.opaque() << " recv_buffer=" << recv_buffer.opaque() << " count=" << count << " dtype=" << primitive_util::LowercasePrimitiveTypeName(dtype) - << " stream=" << AsHipStream(stream); - - ABSL_ASSIGN_OR_RETURN(auto dt, ToMoriDataType(dtype)); - ABSL_ASSIGN_OR_RETURN(auto op, ToMoriReduceOp(kind)); - return se::gpu::ToStatus( - facade_->RunReduceScatter(send_buffer.opaque(), recv_buffer.opaque(), - count, dt, op, AsHipStream(stream))); -} - -absl::Status MoriCommunicator::LaunchAllToAll( - absl::InlinedVector send_buffers, - absl::InlinedVector recv_buffers, - PrimitiveType dtype, size_t count, const Executor& executor) { - CHECK_CANCELLED() - ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - - auto format_addr = [](std::string* out, se::DeviceAddressBase buf) { - absl::StrAppendFormat(out, "%p", buf.opaque()); - }; - VLOG(3) << absl::StreamFormat( - "Launch MORI AllToAll operation; send_buffers=[%s]; recv_buffers=[%s]; " - "dtype=%s; count=%d; stream=%p", - absl::StrJoin(send_buffers, ", ", format_addr), - absl::StrJoin(recv_buffers, ", ", format_addr), - primitive_util::LowercasePrimitiveTypeName(dtype), count, - AsHipStream(stream)); - - if (send_buffers.size() != recv_buffers.size() || - send_buffers.size() != static_cast(num_ranks_)) { - return InvalidArgument( - "Number of send/recv buffers and number of ranks mismatch"); - } - - CollectivesFacade::AddressVector addrs; - addrs.reserve(num_ranks_); - for (int p = 0; p < num_ranks_; ++p) { - addrs.emplace_back(send_buffers[p].opaque(), recv_buffers[p].opaque()); - } - return se::gpu::ToStatus(facade_->RunAllToAll( - addrs, ToMoriByteCount(dtype, count), AsHipStream(stream))); -} - -absl::Status MoriCommunicator::LaunchSend(se::DeviceAddressBase send_buffer, - PrimitiveType dtype, size_t count, - RankId peer, - const Executor& executor) { - CHECK_CANCELLED() - ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - VLOG(3) << absl::StreamFormat( - "Launch MORI Send operation; send_buffer=%p; dtype=%s; count=%d; " - "peer=%d; stream=%p", - send_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), - count, peer.value(), AsHipStream(stream)); - return se::gpu::ToStatus( - facade_->RunSend(send_buffer.opaque(), ToMoriByteCount(dtype, count), - static_cast(peer.value()), AsHipStream(stream))); -} - -absl::Status MoriCommunicator::LaunchRecv(se::DeviceAddressBase recv_buffer, - PrimitiveType dtype, size_t count, - RankId peer, - const Executor& executor) { - CHECK_CANCELLED() - ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - VLOG(3) << absl::StreamFormat( - "Launch MORI Recv operation; recv_buffer=%p; dtype=%s; count=%d; " - "peer=%d; stream=%p", - recv_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), - count, peer.value(), AsHipStream(stream)); - return se::gpu::ToStatus( - facade_->RunRecv(recv_buffer.opaque(), ToMoriByteCount(dtype, count), - static_cast(peer.value()), AsHipStream(stream))); + << " stream=" << AsRocmStream(stream); + return absl::UnimplementedError("Not implemented"); } absl::Status MoriCommunicator::LaunchCollectivePermute( @@ -402,11 +275,13 @@ absl::Status MoriCommunicator::LaunchCollectivePermute( absl::Span target_ranks, const Executor& executor) { CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + size_t bytes = ToMoriByteCount(dtype, count); + (void)bytes; auto rank_formatter = [](std::string* out, RankId rank) { absl::StrAppendFormat(out, "%d", rank.value()); }; VLOG(3) << absl::StreamFormat( - "[%d] Launch CollectivePermute: send_buffer=%p; " + "[%d] Launch MORI CollectivePermute operation; send_buffer=%p; " "recv_buffer=%p; dtype=%s; source_rank=%s; target_[ranks=%s]; count=%d; " "stream=%p", stream->parent()->device_ordinal(), send_buffer.opaque(), @@ -414,15 +289,41 @@ absl::Status MoriCommunicator::LaunchCollectivePermute( source_rank ? absl::StrCat(source_rank->value()) : "", absl::StrJoin(target_ranks, ", ", rank_formatter), count, stream); - std::vector dstPes; - dstPes.reserve(target_ranks.size()); - for (RankId rank : target_ranks) { - dstPes.push_back(static_cast(rank.value())); - } - const int srcPe = source_rank ? static_cast(source_rank->value()) : -1; - return se::gpu::ToStatus(facade_->RunCollectivePermute( - send_buffer.opaque(), recv_buffer.opaque(), ToMoriByteCount(dtype, count), - srcPe, dstPes, AsHipStream(stream))); + return absl::UnimplementedError("Not implemented"); +} + +// Performs point-to-point communication between two ranks using MORI. +// Send: launches a single GPU kernel that copies data to the peer via P2P +// and sets a completion flag on the peer. +// Recv: launches a single-thread GPU kernel that waits for the flag. +absl::Status MoriCommunicator::P2P(P2PType p2p_type, PrimitiveType dtype, + se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, + size_t count, RankId peer, + const Executor& executor) { + const char* stype = (p2p_type == P2PType::Send ? " Send" : " Recv"); + VLOG(1) << CurrentRank().value() << stype << " to " << peer.value() + << " count " << count << " MORI communicator: " << ToString(); + CHECK_CANCELLED() + + void* source_ptr = send_buffer.opaque(); + void* dest_ptr = recv_buffer.opaque(); + + ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + auto gpu_stream = AsRocmStream(stream); + size_t bytes = ToMoriByteCount(dtype, count); + int res = 0; + (void)bytes; + (void)res; + (void)gpu_stream; + (void)source_ptr; + (void)dest_ptr; + (void)peer; + (void)stream; + (void)dtype; + (void)count; + (void)p2p_type; + return absl::UnimplementedError("Not implemented"); } Future<> MoriCommunicator::GroupExecute( @@ -438,16 +339,19 @@ absl::Status MoriCommunicator::GroupLaunch( } absl::Status MoriCommunicator::Quiet(const Executor& executor) { - VLOG(1) << "Quiet: " << ToString(); + VLOG(1) << "Quiet MORI communicator: " << ToString(); CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - return se::gpu::ToStatus(facade_->RunQuiet(AsHipStream(stream))); + auto gpu_stream = AsRocmStream(stream); + (void)gpu_stream; + return absl::UnimplementedError("Not implemented"); } absl::Status MoriCommunicator::Fence() { - VLOG(1) << "Fence: " << ToString(); + VLOG(1) << "Fence MORI communicator: " << ToString(); CHECK_CANCELLED() - return se::gpu::ToStatus(facade_->RunFence()); + // rocm_mori_fence(); + return absl::UnimplementedError("Not implemented"); } absl::Status MoriCommunicator::PollUntilDone() const { diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h index 61b7f14a12c394..ac4b2ffb0e88b7 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h +++ b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h @@ -25,7 +25,6 @@ limitations under the License. #include "absl/functional/function_ref.h" #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/cancellation_token.h" @@ -33,44 +32,20 @@ limitations under the License. #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/reduction_kind.h" -#include "xla/core/collectives/symmetric_memory.h" #include "xla/future.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/stream.h" #include "xla/xla_data.pb.h" -namespace mori::collective { -class CollectivesFacade; -} // namespace mori::collective - namespace xla::gpu { class MoriCollectives; -// Dummy symmetric memory for the MORI backend. MORI collective buffers are -// allocated directly from the symmetric shmem heap, so the local device -// address doubles as the symmetric handle and no separate registration is -// required. This simply returns the address it was created with. -class MoriSymmetricMemory : public SymmetricMemory { - public: - explicit MoriSymmetricMemory(se::DeviceAddressBase addr) : addr_(addr) {} - - se::DeviceAddressBase addr() const final { return addr_; } - - std::string ToString() const final { - return absl::StrFormat("MoriSymmetricMemory(addr=%p, size=%d)", - addr_.opaque(), addr_.size()); - } - - PackedKernelArg PackKernelArg() const final { return addr_.opaque(); } - - private: - se::DeviceAddressBase addr_; -}; - // XLA collectives communicator wrapping a MORI communicator. class MoriCommunicator : public GpuCommunicator { public: + constexpr static uint32_t kMaxTeams = 24; + friend class MoriCollectives; ~MoriCommunicator() override; @@ -88,14 +63,6 @@ class MoriCommunicator : public GpuCommunicator { absl::StatusOr NumRanks() const final; absl::StatusOr CurrentRank() final; - absl::StatusOr> CreateSymmetricMemory( - se::DeviceAddressBase addr) final { - // Dummy implementation: MORI buffers are already allocated from the - // symmetric shmem heap, so the local device address is the symmetric - // handle. Just wrap and return it unchanged. - return std::make_unique(addr); - } - absl::Status Barrier(const Executor& executor) final; Future<> GroupExecute(absl::AnyInvocable group) final; @@ -131,9 +98,21 @@ class MoriCommunicator : public GpuCommunicator { const Executor& executor) final; Future<> Send(se::DeviceAddressBase send_buffer, PrimitiveType dtype, - size_t count, RankId peer, const Executor& executor) final; + size_t count, RankId peer, const Executor& executor) final { + return absl::UnimplementedError("Not implemented"); + } Future<> Recv(se::DeviceAddressBase recv_buffer, PrimitiveType dtype, + size_t count, RankId peer, const Executor& executor) final { + return absl::UnimplementedError("Not implemented"); + } + + Future<> Send(se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, PrimitiveType dtype, + size_t count, RankId peer, const Executor& executor) final; + + Future<> Recv(se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) final; // Polls the communicator until any pending non-blocking operations are done @@ -170,7 +149,9 @@ class MoriCommunicator : public GpuCommunicator { absl::Status LaunchAllToAll( absl::InlinedVector send_buffers, absl::InlinedVector recv_buffers, - PrimitiveType dtype, size_t count, const Executor& executor) final; + PrimitiveType dtype, size_t count, const Executor& executor) final { + return absl::UnimplementedError("Not implemented"); + } absl::Status LaunchCollectivePermute(se::DeviceAddressBase send_buffer, se::DeviceAddressBase recv_buffer, @@ -181,11 +162,15 @@ class MoriCommunicator : public GpuCommunicator { absl::Status LaunchSend(se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, - const Executor& executor) final; + const Executor& executor) final { + return absl::UnimplementedError("Not implemented"); + } absl::Status LaunchRecv(se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, RankId peer, - const Executor& executor) final; + const Executor& executor) final { + return absl::UnimplementedError("Not implemented"); + } absl::Status Quiet(const Executor& executor) final; @@ -201,16 +186,22 @@ class MoriCommunicator : public GpuCommunicator { std::shared_ptr cancel) : collectives_(coll), cancel_(std::move(cancel)) {} + enum class P2PType : int32_t { Send, Recv }; + + absl::Status P2P(P2PType p2p_type, PrimitiveType type, + se::DeviceAddressBase recv_buffer, + se::DeviceAddressBase send_buffer, size_t count, RankId peer, + const Executor& executor); + + static absl::StatusOr ToStream(const Executor& executor); + MoriCollectives* collectives_; // Parent MoriCollectives instance // This communicator's participant set (NOT the global MORI clique). `rank_` - // is this rank within the collective, `num_ranks_` the participant count. + // is this rank within the collective, `num_ranks_` the participant count, and + // `rank_to_pe_dev_` a device array mapping collective rank -> global MORI PE. int rank_ = 0; int num_ranks_ = 0; - // Owns this communicator's staging buffer + group counters (created in - // Create(), freed by the facade dtor before ShmemFinalize). Header-only - // facade. - std::unique_ptr<::mori::collective::CollectivesFacade> facade_; // Should all pending collectives cancel? std::shared_ptr cancel_; bool aborted_ = false; // Has Abort() been called? diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc deleted file mode 100644 index 76b616a1942d22..00000000000000 --- a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc +++ /dev/null @@ -1,19 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -// This is the single device translation unit for the MORI XLA collectives. It -// is compiled as HIP and defines MORI_KERNELS_IMPL before including the facade, -// so the facade's device path (kernels + non-templated Run* definitions) is -// compiled here exactly once. The host mori_communicator.cc includes the same -// header without MORI_KERNELS_IMPL (decl-only) and links against these symbols. -#define MORI_KERNELS_IMPL -#include "xla/backends/gpu/collectives/mori_kernels.h" diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h deleted file mode 100644 index bff6bc1860d5e8..00000000000000 --- a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h +++ /dev/null @@ -1,27 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ -#define XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ - -#include -#include - -// The CollectivesFacade owns the per-device staging + Run* entry points, which -// are non-templated and take mori::collective::DataType / ReduceOpKind enums. -// Host includers (mori_communicator.cc) see decl-only Run* methods; the device -// TU (mori_kernels.cu.cc, compiled as HIP with MORI_KERNELS_IMPL) pulls in the -// full device path and emits the definitions that resolve the host's -// references. -#include "xla/backends/gpu/collectives/mori_stub.h" - -#endif // XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_stub.h b/third_party/xla/xla/backends/gpu/collectives/mori_stub.h index 2a18155d71dc17..267cdf3a35fbe2 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_stub.h +++ b/third_party/xla/xla/backends/gpu/collectives/mori_stub.h @@ -16,14 +16,9 @@ limitations under the License. #ifndef XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ #define XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ -#include - #include #include #include -#include -#include -#include // Inert stand-in for the subset of the MORI shmem host API used by the MORI // collectives/communicator backbone. These placeholders let the backbone @@ -77,74 +72,4 @@ inline void ShmemFree(void* /*ptr*/) {} } // namespace shmem } // namespace mori -namespace mori { -namespace collective { - -// Element type + reduction op enums mirror the real facade's non-templated API -// (mori/collective/collectives_facade.hpp), so the communicator's enum dispatch -// compiles against either the stub or the real facade. -enum class DataType { - F8E5M2, - F8E4M3FN, - F16, - BF16, - S8, - U8, - S32, - U32, - S64, - U64, - F32, - F64 -}; -enum class ReduceOpKind { SUM, PRODUCT, MIN, MAX }; - -// Inert stand-in for the real MORI CollectivesFacade. Header-only, all Run* are -// no-ops returning hipSuccess. Lets the collectives/communicator wiring compile -// and link without @roc_mori. -class CollectivesFacade { - CollectivesFacade() = default; - - public: - using AddressVector = std::vector>; - - CollectivesFacade(const CollectivesFacade&) = delete; - CollectivesFacade& operator=(const CollectivesFacade&) = delete; - - static std::unique_ptr Create(int /*myPe*/, int /*nPes*/, - size_t /*maxStagingBytes*/) { - return std::unique_ptr(new CollectivesFacade()); - } - ~CollectivesFacade() = default; - - hipError_t RunReduceScatter(const void*, void*, size_t, DataType, - ReduceOpKind, hipStream_t) { - return hipSuccess; - } - hipError_t RunAllReduce(const void*, void*, size_t, DataType, ReduceOpKind, - hipStream_t) { - return hipSuccess; - } - hipError_t RunAllGather(const void*, void*, size_t, hipStream_t) { - return hipSuccess; - } - hipError_t RunAllToAll(const AddressVector&, size_t, hipStream_t) { - return hipSuccess; - } - hipError_t RunBarrier(hipStream_t) { return hipSuccess; } - hipError_t RunSend(const void*, size_t, int, hipStream_t) { - return hipSuccess; - } - hipError_t RunRecv(void*, size_t, int, hipStream_t) { return hipSuccess; } - hipError_t RunCollectivePermute(const void*, void*, size_t, int, - const std::vector&, hipStream_t) { - return hipSuccess; - } - hipError_t RunQuiet(hipStream_t) { return hipSuccess; } - hipError_t RunFence() { return hipSuccess; } -}; - -} // namespace collective -} // namespace mori - #endif // XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ From 0061623663cdce96884fdca5fe4c0fa0953c06b0 Mon Sep 17 00:00:00 2001 From: Yang Chen Date: Tue, 1 Sep 2026 17:52:46 -0700 Subject: [PATCH 6/8] Add log to confirm tf.data service num_clients is set correctly. PiperOrigin-RevId: 974806669 --- tensorflow/core/data/service/dispatcher_state.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/data/service/dispatcher_state.cc b/tensorflow/core/data/service/dispatcher_state.cc index 92f01199ee2200..1b846d46e3dde0 100644 --- a/tensorflow/core/data/service/dispatcher_state.cc +++ b/tensorflow/core/data/service/dispatcher_state.cc @@ -203,6 +203,9 @@ void DispatcherState::AcquireIterationClient( iteration->num_clients++; next_available_iteration_client_id_ = std::max(next_available_iteration_client_id_, iteration_client_id + 1); + VLOG(3) << "Acquired iteration client for iteration " + << iteration->iteration_id + << " num_clients: " << iteration->num_clients; } void DispatcherState::ReleaseIterationClient( @@ -215,6 +218,9 @@ void DispatcherState::ReleaseIterationClient( DCHECK_GE(iteration->num_clients, 0); iteration->last_client_released_micros = release_iteration_client.time_micros(); + VLOG(3) << "Released iteration client for iteration " + << iteration->iteration_id + << " num_clients: " << iteration->num_clients; iterations_for_client_ids_.erase(iteration_client_id); } From 2f927e4b46f7db6d94f34e91a09ccb3101c2f41a Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Tue, 1 Sep 2026 19:52:35 -0700 Subject: [PATCH 7/8] Scale self-adjoint eig input by max element to fix numerical instability Self-adjoint eigendecomposition can suffer from numerical instability or underflow/overflow during two-sided Jacobi rotations when matrix elements have very large or very small magnitudes. Scale the input matrix by the maximum absolute element before performing eigensystem decomposition, and scale the computed eigenvalues back by the same factor. Properly handle both real and complex component types for TPU and CPU/GPU expander implementations. PiperOrigin-RevId: 974847313 --- .../hlo/builder/lib/self_adjoint_eig_test.cc | 113 ++++++++++++++++++ .../hlo/transforms/expanders/eigh_expander.cc | 58 +++++++-- .../hlo/transforms/expanders/eigh_expander.h | 8 ++ 3 files changed, 169 insertions(+), 10 deletions(-) diff --git a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc index 467445ea097988..5cfdd6f601c784 100644 --- a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc +++ b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc @@ -269,6 +269,119 @@ TEST_F(SelfAdjointEigTest, Wrong_Type_Int) { EXPECT_FALSE(result.w.valid()); } +TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_2x2) { + XlaBuilder builder(TestName()); + float v = 1e20f; + Array2D input{{v, v}, {v, v}}; + std::vector expected{0.0f, 2e20f}; + + XlaOp a; + auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); + auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, + /*tol=*/1e-5, /*sort_eigenvalues=*/true); + Add(result.w, ZerosLike(result.w)); + + ComputeAndCompareR1(&builder, expected, {&a_data}, + ErrorSpec(1e15f, 1e-4f)); +} + +TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_3x3) { + XlaBuilder builder(TestName()); + float v = 1e20f; + Array2D input{{v, v, v}, {v, v, v}, {v, v, v}}; + std::vector expected{0.0f, 0.0f, 3e20f}; + + XlaOp a; + auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); + auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, + /*tol=*/1e-5, /*sort_eigenvalues=*/true); + Add(result.w, ZerosLike(result.w)); + + ComputeAndCompareR1(&builder, expected, {&a_data}, + ErrorSpec(1e15f, 1e-4f)); +} + +TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_Complex_3x3) { + XlaBuilder builder(TestName()); + float v = 1e20f; + Array input = { + {complex64{v, 0.0f}, complex64{v, -v}, complex64{0.0f, 0.0f}}, + {complex64{v, v}, complex64{v, 0.0f}, complex64{0.0f, 0.0f}}, + {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{v, 0.0f}}, + }; + const Literal a_literal = LiteralUtil::CreateFromArray(input); + XlaOp a = Parameter(&builder, 0, a_literal.shape(), "a"); + auto result = SelfAdjointEig(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareLiteral(&builder, LiteralUtil::CreateFromArray(input), + {&a_literal}, ErrorSpec(1e15f, 1e-4f)); +} + +TEST_F(SelfAdjointEigTest, Test_Small_Magnitude_2x2) { + XlaBuilder builder(TestName()); + float v = 1e-20f; + Array2D input{{v, v}, {v, v}}; + std::vector expected{0.0f, 2e-20f}; + + XlaOp a; + auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); + auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, + /*tol=*/1e-5, /*sort_eigenvalues=*/true); + Add(result.w, ZerosLike(result.w)); + + ComputeAndCompareR1(&builder, expected, {&a_data}, + ErrorSpec(1e-25f, 1e-4f)); +} + +TEST_F(SelfAdjointEigTest, Test_Small_Magnitude_Complex_3x3) { + XlaBuilder builder(TestName()); + float v = 1e-20f; + Array input = { + {complex64{v, 0.0f}, complex64{v, -v}, complex64{0.0f, 0.0f}}, + {complex64{v, v}, complex64{v, 0.0f}, complex64{0.0f, 0.0f}}, + {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{v, 0.0f}}, + }; + const Literal a_literal = LiteralUtil::CreateFromArray(input); + XlaOp a = Parameter(&builder, 0, a_literal.shape(), "a"); + auto result = SelfAdjointEig(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareLiteral(&builder, LiteralUtil::CreateFromArray(input), + {&a_literal}, ErrorSpec(1e-25f, 1e-4f)); +} + +TEST_F(SelfAdjointEigTest, Test_Zero_Matrix_2x2) { + XlaBuilder builder(TestName()); + Array2D input{{0.0f, 0.0f}, {0.0f, 0.0f}}; + std::vector expected{0.0f, 0.0f}; + + XlaOp a; + auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); + auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, + /*tol=*/1e-5, /*sort_eigenvalues=*/true); + Add(result.w, ZerosLike(result.w)); + + ComputeAndCompareR1(&builder, expected, {&a_data}, + ErrorSpec(1e-6f, 1e-6f)); +} + +TEST_F(SelfAdjointEigTest, Test_Zero_Matrix_Complex_3x3) { + XlaBuilder builder(TestName()); + Array input = { + {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}}, + {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}}, + {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}}, + }; + const Literal a_literal = LiteralUtil::CreateFromArray(input); + XlaOp a = Parameter(&builder, 0, a_literal.shape(), "a"); + auto result = SelfAdjointEig(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareLiteral(&builder, LiteralUtil::CreateFromArray(input), + {&a_literal}, ErrorSpec(1e-6f, 1e-6f)); +} + Array2D GenerateRandomSymmetricMatrix(int size) { Array2D result{size, size, 0.0}; // TODO(b/128001705): This seed should not be needed but makes the test diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc index 3b6c0debf5a26e..4e395d3f68a926 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc @@ -15,7 +15,6 @@ limitations under the License. #include "xla/hlo/transforms/expanders/eigh_expander.h" -#include #include #include #include @@ -351,7 +350,7 @@ absl::StatusOr> Sweeps( sweep_values, "ApplyRotations", body_builder)); std::vector output(values.size()); output[0] = values[0] + ScalarLike(values[0], 1); - std::copy(sweep_values.begin(), sweep_values.end(), output.begin() + 1); + absl::c_copy(sweep_values, output.begin() + 1); return output; }; return WhileLoopHelper(while_cond_fn, while_body_fn, initial_values, @@ -435,6 +434,47 @@ absl::Status EighExpander::SortByEigenvalues(XlaOp& v, XlaOp& w) { // off_diag_norm = np.sqrt(frobenius_norm - diag_norm) * np.sqrt( // frobenius_norm + diag_norm) // return A, V +absl::StatusOr EighExpander::ScaleInputMatrix( + XlaOp a) { + XlaBuilder* builder = a.builder(); + ABSL_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + const int64_t num_dims = a_shape.dimensions().size(); + const int64_t num_batch_dims = num_dims - 2; + PrimitiveType type = a_shape.element_type(); + PrimitiveType real_type = primitive_util::IsComplexType(type) + ? primitive_util::ComplexComponentType(type) + : type; + XlaOp zero_real = Zero(builder, real_type); + XlaOp one_real = One(builder, real_type); + XlaOp abs_a = primitive_util::IsComplexType(type) + ? Max(Abs(Real(a)), Abs(Imag(a))) + : Abs(a); + XlaOp a_max = + Reduce(abs_a, zero_real, CreateScalarMaxComputation(real_type, builder), + {num_dims - 2, num_dims - 1}); + XlaOp scale = Select(Eq(a_max, zero_real), one_real, a_max); + + std::vector batch_broadcast_dims(num_batch_dims); + absl::c_iota(batch_broadcast_dims, 0); + + XlaOp scale_a = primitive_util::IsComplexType(type) + ? Complex(scale, ZerosLike(scale)) + : scale; + scale_a = BroadcastInDim(scale_a, a_shape.dimensions(), batch_broadcast_dims); + return ScaledInput{a / scale_a, scale}; +} + +absl::StatusOr EighExpander::RescaleEigenvalues(XlaOp w, XlaOp scale) { + XlaBuilder* builder = w.builder(); + ABSL_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w)); + const int64_t num_batch_dims = w_shape.dimensions().size() - 1; + std::vector batch_broadcast_dims(num_batch_dims); + absl::c_iota(batch_broadcast_dims, 0); + XlaOp scale_w = + BroadcastInDim(scale, w_shape.dimensions(), batch_broadcast_dims); + return w * scale_w; +} + XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, bool sort_eigenvalues) { XlaBuilder* builder = a.builder(); @@ -477,11 +517,11 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, a = Symmetrize(a, lower); + ABSL_ASSIGN_OR_RETURN(ScaledInput scaled_input, ScaleInputMatrix(a)); + a = scaled_input.scaled_matrix; + XlaOp scale = scaled_input.scale; + const int64_t k = CeilOfRatio(n, int64_t{2}); - // tl = A[:n // 2, :n // 2] - // bl = A[n // 2:, :n // 2] - // tr = A[:n // 2, n // 2:] - // br = A[n // 2:, n // 2:] auto tl = SliceInMinorDims(a, {0, 0}, {k, k}); auto bl = SliceInMinorDims(a, {k, 0}, {n, k}); auto tr = SliceInMinorDims(a, {0, k}, {k, n}); @@ -495,10 +535,6 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, config.mutable_dimensions(num_dims - 1)->set_edge_padding_high(1); br = Pad(br, zero, config); } - // v_tl = np.eye(n // 2, dtype=A.dtype) - // v_tr = np.zeros((n // 2, n // 2), A.dtype) - // v_bl = np.zeros((n // 2, n // 2), A.dtype) - // v_br = np.eye(n // 2, dtype=A.dtype) auto v_tl = Broadcast(IdentityMatrix(builder, type, k, k), batch_dims); auto v_br = v_tl; auto v_tr = ZerosLike(v_tl); @@ -537,6 +573,8 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, } v = MaybeConjugate(TransposeInMinorDims(v), true); + ABSL_ASSIGN_OR_RETURN(w, RescaleEigenvalues(w, scale)); + if (sort_eigenvalues) { ABSL_RETURN_IF_ERROR(SortByEigenvalues(v, w)); } diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.h b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.h index 3f47d792183de1..33785175310215 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.h +++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.h @@ -41,6 +41,14 @@ class EighExpander : public OpExpanderPass { virtual XlaOp BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, bool sort_eigenvalues); + struct ScaledInput { + XlaOp scaled_matrix; + XlaOp scale; + }; + + static absl::StatusOr ScaleInputMatrix(XlaOp a); + static absl::StatusOr RescaleEigenvalues(XlaOp w, XlaOp scale); + absl::Status SortByEigenvalues(XlaOp& v, XlaOp& w); private: From 54288edf8781c60d770a7e1dd00ad7bac7586f04 Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Tue, 1 Sep 2026 21:16:08 -0700 Subject: [PATCH 8/8] Add constant legalizations and optimizer fusions in LiteRT converter pipeline - Legalize VHLO quant custom calls during bytecode module loading in slim_model_importer. - Constrain MHLO constant legalization to Arith constants when buildable, falling back to TFLite constants. - Add patterns in optimize_broadcast_like_patterns to fold splat TFL constants and TFL Fill ops into SelectV2. - Add pattern and CanFuseSumMulToMean helper to fuse Sum followed by 1/N multiplication into Mean. PiperOrigin-RevId: 974875739 --- tensorflow/compiler/mlir/lite/BUILD | 2 +- .../tflite_legalize_hlo_patterns.td | 11 +++++ .../compiler/mlir/lite/tests/optimize.mlir | 27 +++++++++++ .../optimize_broadcast_like_patterns.td | 48 +++++++++++++++++++ .../mlir/lite/transforms/optimize_pass.cc | 44 ++++++++++++++++- .../mlir/lite/transforms/optimize_patterns.td | 40 ++++++++++++++++ 6 files changed, 170 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index 26fbb92a7e3fb1..435b63a90d65b5 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -1390,6 +1390,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_set", "@llvm-project//llvm:Support", "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:Dialect", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", @@ -1827,7 +1828,6 @@ cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", - "@eigen_archive//:eigen3", "@flatbuffers", "@llvm-project//llvm:Support", "@llvm-project//mlir:ArithDialect", diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td index 2f030d07c6a099..db16a99353ba85 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td @@ -31,9 +31,20 @@ def CreateTFLCastToInt32Op : NativeCodeCall< def ConstDefaultResultAccuracyAttr : ConstantAttr; +def IsBuildableWithArithConstant : Constraint< + CPred<"::mlir::arith::ConstantOp::isBuildableWith($0, $1.getType())"> +>; + def : Pat< (MHLO_ConstantOp:$output $value), (Arith_ConstantOp $value), + [(IsBuildableWithArithConstant $value, $output)], + [], + (addBenefit 10)>; + +def : Pat< + (MHLO_ConstantOp:$output $value), + (TFL_ConstOp $value), [(TFL_TensorOf<[AnyType]> $output)]>; diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir index 71768d8b8bf87f..c6e18be30ae02d 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir @@ -5165,3 +5165,30 @@ func.func @Fuse4DResourceAddIntoDepthwiseConv2D(%arg0: tensor<1x32x32x4xf32>, %a // CHECK: %[[DW_CONV:.*]] = "tfl.depthwise_conv_2d"(%arg0, %arg1, %[[NEW_BIAS]]) // CHECK: return %[[DW_CONV]] } + +// CHECK-LABEL: @fuse_sum_mul_to_mean +func.func @fuse_sum_mul_to_mean(%arg0: tensor<2x3x4xf32>) -> tensor<2x1x4xf32> { + %cst_axes = arith.constant dense<1> : tensor<1xi32> + %cst_factor = arith.constant dense<0.333333333> : tensor<1xf32> + %0 = "tfl.sum"(%arg0, %cst_axes) <{keep_dims = true}> : (tensor<2x3x4xf32>, tensor<1xi32>) -> tensor<2x1x4xf32> + %1 = "tfl.mul"(%0, %cst_factor) <{fused_activation_function = "NONE"}> : (tensor<2x1x4xf32>, tensor<1xf32>) -> tensor<2x1x4xf32> + func.return %1 : tensor<2x1x4xf32> + + // CHECK-NOT: tfl.sum + // CHECK-NOT: tfl.mul + // CHECK: "tfl.mean"(%arg0, %{{.*}}) <{keep_dims = true}> +} + +// CHECK-LABEL: @do_not_fuse_sum_mul_with_arbitrary_factor +func.func @do_not_fuse_sum_mul_with_arbitrary_factor(%arg0: tensor<2x3x4xf32>) -> tensor<2x1x4xf32> { + %cst_axes = arith.constant dense<1> : tensor<1xi32> + %cst_factor = arith.constant dense<8.000000e-01> : tensor<1xf32> + %0 = "tfl.sum"(%arg0, %cst_axes) <{keep_dims = true}> : (tensor<2x3x4xf32>, tensor<1xi32>) -> tensor<2x1x4xf32> + %1 = "tfl.mul"(%0, %cst_factor) <{fused_activation_function = "NONE"}> : (tensor<2x1x4xf32>, tensor<1xf32>) -> tensor<2x1x4xf32> + func.return %1 : tensor<2x1x4xf32> + + // CHECK: "tfl.sum" + // CHECK: tfl.mul + // CHECK-NOT: "tfl.mean" +} + diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_patterns.td index 945c67090f08fd..bee713efc1a36e 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_patterns.td +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_patterns.td @@ -58,6 +58,33 @@ multiclass FuseSplatConstIntoSelectOp { [(HasRankAtLeast<2> $constant_attr), (OperandsBroadcastToOutputType $input1, $input2, $result), (HasRankAtMost<5> $constant_value)]>; + + def FuseSplatTflConstLhsInto#SelectOp : Pat< + (SelectOp:$result + AnyStaticShapeTensor:$input1, + (TFL_ConstOp:$constant_value SplatElementsAttr:$constant_attr), + AnyStaticShapeTensor:$input2), + (TFL_SelectV2Op + $input1, + (TFL_ConstOp (GetScalarElementsAttrFromSplat $constant_attr)), + $input2), + // Check if condition or rhs will promote the required broadcasting. + [(HasRankAtLeast<2> $constant_attr), + (OperandsBroadcastToOutputType $input1, $input2, $result), + (HasRankAtMost<5> $constant_value)]>; + + def FuseSplatTflConstRhsInto#SelectOp : Pat< + (SelectOp:$result + AnyStaticShapeTensor:$input1, + AnyStaticShapeTensor:$input2, + (TFL_ConstOp:$constant_value SplatElementsAttr:$constant_attr)), + (TFL_SelectV2Op + $input1, $input2, + (TFL_ConstOp (GetScalarElementsAttrFromSplat $constant_attr))), + // Check if condition or lhs will promote the required broadcasting. + [(HasRankAtLeast<2> $constant_attr), + (OperandsBroadcastToOutputType $input1, $input2, $result), + (HasRankAtMost<5> $constant_value)]>; } // Pattern for skipping FillOp if it is mainly for broadcasting and the @@ -83,6 +110,27 @@ multiclass FuseFillOpBroadcastIntoFollowingSelectOp { [(OperandsBroadcastToOutputType $input1, $input2, $result), (HasRankAtMost<5> $result), (IsRankLessThanEqualTo $fill_output, $result)]>; + + def FoldTflFillOpIntoSelectOpRHS#SelectOp : Pat< + (SelectOp:$result + AnyStaticShapeTensor:$input1, + AnyStaticShapeTensor:$input2, + (TFL_FillOp:$fill_output $fill_dims, + (TFL_ConstOp:$fill_value $val))), + (TFL_SelectV2Op $input1, $input2, $fill_value), + [(OperandsBroadcastToOutputType $input1, $input2, $result), + (HasRankAtMost<5> $result), + (IsRankLessThanEqualTo $fill_output, $result)]>; + + def FoldTflFillOpIntoSelectOpLHS#SelectOp : Pat< + (SelectOp:$result + AnyStaticShapeTensor:$input1, + (TFL_FillOp:$fill_output $fill_dims, (TFL_ConstOp:$fill_value $val)), + AnyStaticShapeTensor:$input2), + (TFL_SelectV2Op $input1, $fill_value, $input2), + [(OperandsBroadcastToOutputType $input1, $input2, $result), + (HasRankAtMost<5> $result), + (IsRankLessThanEqualTo $fill_output, $result)]>; } multiclass FuseBroadcastToIntoSelectOp { diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc index ac6744b5c4fa32..49e73d98ba4569 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc @@ -21,11 +21,12 @@ limitations under the License. #include #include #include +#include #include #include +#include #include #include -#include #include #include #include @@ -42,12 +43,14 @@ limitations under the License. #include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project +#include "mlir/Dialect/Traits.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Matchers.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project @@ -572,6 +575,45 @@ TypeAttr RescaleQtype(Type input, Attribute factor) { return RescaleQuantizedType(input, factor); } +// Returns true if the multiplication by factor (1/N) following a sum reduction +// along axes can be fused into TFL::MeanOp. +static bool CanFuseSumMulToMean(Value input, Attribute axes, Attribute factor) { + RankedTensorType input_type = + mlir::dyn_cast_or_null(input.getType()); + if (!input_type || !input_type.hasStaticShape()) return false; + + auto dense_factor = mlir::dyn_cast_or_null(factor); + if (!dense_factor || dense_factor.getNumElements() != 1) return false; + if (!mlir::isa(dense_factor.getElementType())) return false; + + float factor_val = + (*dense_factor.getValues().begin()).convertToFloat(); + if (factor_val <= 0.0f) return false; + + auto dense_axes = mlir::dyn_cast_or_null(axes); + if (!dense_axes || dense_axes.empty()) return false; + + int64_t rank = input_type.getRank(); + auto shape = input_type.getShape(); + int64_t reduction_elements = 1; + + llvm::SmallSet unique_axes; + for (const APInt& val : dense_axes.getValues()) { + int64_t axis = val.getSExtValue(); + if (axis < 0) axis += rank; + if (axis < 0 || axis >= rank) return false; + if (unique_axes.insert(axis).second) { + reduction_elements *= shape[axis]; + } + } + + if (reduction_elements <= 0) return false; + + float expected_factor = 1.0f / static_cast(reduction_elements); + float diff = std::abs(factor_val - expected_factor); + return diff <= 1e-4f * expected_factor; +} + // Returns `true` if reducing `axes` in `input` with `keep_dims=true` results // in the specified `shape` and `false` otherwise. static bool ShapeMatchesReduceWithKeepAxes(Value input, diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td index b7db418465526b..d1940e9aee4fec 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td @@ -1076,6 +1076,46 @@ foreach ReduceOp = [TFL_MeanOp, TFL_ReduceMaxOp, TFL_ReduceMinOp, (HasOneUse $reduce)]>; } +def CanFuseSumMulToMean : Constraint>; + +// Fuse Sum + Mul(1/N) -> Mean +def FuseSumMulRhsIntoMean : Pat< + (TFL_MulOp + (TFL_SumOp:$sum $input, (Arith_ConstantOp ElementsAttr:$axes), $keep_dims), + (Arith_ConstantOp ElementsAttr:$factor), + TFL_AF_None), + (TFL_MeanOp $input, (Arith_ConstantOp $axes), $keep_dims), + [(CanFuseSumMulToMean $input, $axes, $factor), + (HasOneUse $sum)]>; + +def FuseSumMulLhsIntoMean : Pat< + (TFL_MulOp + (Arith_ConstantOp ElementsAttr:$factor), + (TFL_SumOp:$sum $input, (Arith_ConstantOp ElementsAttr:$axes), $keep_dims), + TFL_AF_None), + (TFL_MeanOp $input, (Arith_ConstantOp $axes), $keep_dims), + [(CanFuseSumMulToMean $input, $axes, $factor), + (HasOneUse $sum)]>; + +def FuseTflSumMulRhsIntoMean : Pat< + (TFL_MulOp + (TFL_SumOp:$sum $input, (TFL_ConstOp ElementsAttr:$axes), $keep_dims), + (TFL_ConstOp ElementsAttr:$factor), + TFL_AF_None), + (TFL_MeanOp $input, (TFL_ConstOp $axes), $keep_dims), + [(CanFuseSumMulToMean $input, $axes, $factor), + (HasOneUse $sum)]>; + +def FuseTflSumMulLhsIntoMean : Pat< + (TFL_MulOp + (TFL_ConstOp ElementsAttr:$factor), + (TFL_SumOp:$sum $input, (TFL_ConstOp ElementsAttr:$axes), $keep_dims), + TFL_AF_None), + (TFL_MeanOp $input, (TFL_ConstOp $axes), $keep_dims), + [(CanFuseSumMulToMean $input, $axes, $factor), + (HasOneUse $sum)]>; + def IsSame : Constraint>; def HasTwoUse : Constraint