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/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/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 : Constraintnum_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); } 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. 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", }), }, 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: "" 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_ 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: 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