diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index d1fbed6a8783ed..efc6869f6235de 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -206,12 +206,6 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config, flags.sanitize_abilists_dataflow, ',', absl::SkipEmpty())); } - if (flags.sanitize_memory || flags.sanitize_memory_track_origins > 0) { - aot_opts.set_sanitize_memory(true); - aot_opts.set_sanitize_memory_track_origins( - flags.sanitize_memory_track_origins); - } - TF_RETURN_IF_ERROR( ConfigureKernelNamingConvention(aot_opts, computation, flags.cpp_class)); diff --git a/tensorflow/compiler/aot/flags.cc b/tensorflow/compiler/aot/flags.cc index 6555d22d05c325..567426f53c7631 100644 --- a/tensorflow/compiler/aot/flags.cc +++ b/tensorflow/compiler/aot/flags.cc @@ -86,12 +86,6 @@ void AppendMainFlags(std::vector* flag_list, MainFlags* flags) { "Enable DataFlow Sanitizer pass."}, {"sanitize_abilists_dataflow", &flags->sanitize_abilists_dataflow, "Comma separated list of ABIList file paths."}, - {"sanitize_memory", &flags->sanitize_memory, - "Enable Memory Sanitizer pass."}, - {"sanitize_memory_track_origins", &flags->sanitize_memory_track_origins, - "Controls MSan track origins level (0=disabled, 1=without store " - "history, 2=with store history). Setting to >0 implies " - "--sanitize_memory."}, {"gen_name_to_index", &flags->gen_name_to_index, "Generate name-to-index data for Lookup{Arg,Result}Index methods."}, {"gen_program_shape", &flags->gen_program_shape, diff --git a/tensorflow/compiler/aot/flags.h b/tensorflow/compiler/aot/flags.h index f71fb80e0973f4..5d0f93f7d67b88 100644 --- a/tensorflow/compiler/aot/flags.h +++ b/tensorflow/compiler/aot/flags.h @@ -48,8 +48,6 @@ struct MainFlags { // Sanitizer pass options bool sanitize_dataflow = false; std::string sanitize_abilists_dataflow; - bool sanitize_memory = false; - int32_t sanitize_memory_track_origins = 0; // C++ codegen options bool gen_name_to_index = false; diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 83c41ba27c2d20..911c917350d775 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -358,7 +358,6 @@ def _tf_library( "@xla//xla/backends/cpu/runtime:sort_lib", "@xla//xla/backends/cpu/runtime:topk_lib", "@xla//xla/backends/cpu/runtime:convolution_lib", - "@xla//xla/backends/cpu/runtime:msan_emulated_tls", "@xla//xla/service/cpu:runtime_matmul", "@xla//xla/service/cpu:runtime_single_threaded_matmul", "@eigen_archive//:eigen3", diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc index a7d7a2b6e1dd6a..7d34e71a38aa15 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc @@ -56,6 +56,7 @@ limitations under the License. #include "flatbuffers/vector.h" // from @flatbuffers #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/Hashing.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" @@ -64,6 +65,7 @@ limitations under the License. #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/SwapByteOrder.h" #include "llvm/Support/raw_ostream.h" +#include "llvm/Support/xxhash.h" #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 @@ -559,10 +561,10 @@ struct SignatureDefData { // Note, we are using maps here to make order deterministic // for easily testing only. - // Inputs defined in the signature def mapped to tensor names. - std::map inputs; - // Outputs defined in the signature def mapped to tensor names. - std::map outputs; + // Inputs defined in the signature def mapped to tensor index. + std::map inputs; + // Outputs defined in the signature def mapped to tensor index. + std::map outputs; // Signature key. std::string signature_key; // Subgraph index. @@ -824,10 +826,10 @@ class Translator { // Returns list of offsets for the passed 'items' in TensorMap structure // inside the flatbuffer. - // 'items' is a map from tensor name in signatureDef to tensor name in - // the subgraph, specified by the 'subgraph_index' argument. + // 'items' is a map from tensor name in signatureDef to tensor index in + // the subgraph. std::vector> GetList( - int subgraph_index, const std::map& items); + const std::map& items); // Uses the tf.entry_function attribute (if set) to initialize the op to name // mapping. @@ -950,6 +952,10 @@ class Translator { absl::flat_hash_map> tensor_index_map_; + // Stores input and output tensor indices for each subgraph. + std::vector> subgraph_inputs_; + std::vector> subgraph_outputs_; + // Maps op name to index of the corresponding OperatorCode in opcodes_ vector. absl::flat_hash_map opcode_index_map_; std::vector> opcodes_; @@ -1107,6 +1113,93 @@ Translator::BuildExternalBuffer(mlir::Value value, return external_buffer; } +static mlir::AsmResourceBlob* GetBlob( + mlir::DenseResourceElementsAttr resource_attr) { + mlir::AsmResourceBlob* blob = resource_attr.getRawHandle().getBlob(); + if (!blob && resource_attr.getRawHandle().getResource()) { + blob = resource_attr.getRawHandle().getResource()->getBlob(); + } + return blob; +} + +static uint64_t GetPhysicalBufferHash(mlir::ElementsAttr attr) { + if (auto resource_attr = + mlir::dyn_cast(attr)) { + mlir::AsmResourceBlob* blob = GetBlob(resource_attr); + uint64_t h = 0; + if (blob && !blob->getData().empty()) { + h = llvm::xxh3_64bits( + reinterpret_cast(blob->getData().data()), + blob->getData().size()); + } else { + h = llvm::hash_value(resource_attr.getRawHandle().getKey().str()); + } + return llvm::hash_combine(h, mlir::hash_value(resource_attr.getType())); + } + return mlir::hash_value(attr); +} + +static int GetLowBitWidth(tflite::TensorType type) { + switch (type) { + case tflite::TensorType_INT4: + case tflite::TensorType_UINT4: + return 4; + case tflite::TensorType_INT2: + return 2; + default: + return 0; + } +} + +static absl::Status PackLowBitElementsAttr( + mlir::Attribute attr, int bit_width, + absl::FunctionRef apply) { + std::optional raw_data; + size_t num_elements = 0; + + if (auto res_attr = mlir::dyn_cast(attr)) { + if (auto* blob = GetBlob(res_attr); blob && !blob->getData().empty()) { + raw_data = absl::string_view( + reinterpret_cast(blob->getData().data()), + blob->getData().size()); + num_elements = res_attr.getNumElements(); + } + } else if (auto dense_attr = mlir::dyn_cast(attr)) { + if (!dense_attr.isSplat() && !dense_attr.getRawData().empty()) { + raw_data = absl::string_view(dense_attr.getRawData().data(), + dense_attr.getRawData().size()); + num_elements = dense_attr.getNumElements(); + } + } + + // 1. Raw byte buffer path (Resource blobs or DenseElementsAttr) + if (raw_data.has_value()) { + if (raw_data->size() == num_elements) { + if (bit_width == 4) { + return tflite::StreamPackLowBitValues8Bit(*raw_data, + apply); + } else if (bit_width == 2) { + return tflite::StreamPackLowBitValues8Bit(*raw_data, + apply); + } + } + return apply(*raw_data); + } + + // 2. Fallback for splat DenseElementsAttr (e.g., dense<3> : tensor<128xi4>) + if (auto dense_attr = mlir::dyn_cast(attr)) { + if (bit_width == 4) { + return tflite::StreamPackLowBitValues( + dense_attr.getValues(), apply); + } else if (bit_width == 2) { + return tflite::StreamPackLowBitValues( + dense_attr.getValues(), apply); + } + } + + return apply(absl::string_view()); +} + std::optional> Translator::BuildBuffer( mlir::Value value, bool can_be_deduplicated, int& index) { can_be_deduplicated = can_be_deduplicated && !disable_buffer_deduping_; @@ -1159,37 +1252,14 @@ std::optional> Translator::BuildBuffer( GetTFLiteType(type.getElementType()).value(); // Default appliers - if (tflite_element_type == tflite::TensorType_INT4 || - tflite_element_type == tflite::TensorType_UINT4 || - tflite_element_type == tflite::TensorType_INT2) { + int low_bit_width = GetLowBitWidth(tflite_element_type); + if (low_bit_width > 0) { applier = - [tflite_element_type]( + [low_bit_width]( const std::pair& attr_and_inst, auto apply) { - auto attr = mlir::cast(attr_and_inst.first); - bool is_8bit_raw_data = - !attr.isSplat() && - attr.getNumElements() == attr.getRawData().size(); - bool is_4bit_data = tflite_element_type == tflite::TensorType_INT4 || - tflite_element_type == tflite::TensorType_UINT4; - - if (is_8bit_raw_data) { - if (is_4bit_data) { - return tflite::StreamPackLowBitValues8Bit( - attr.getRawData(), apply); - } else { - return tflite::StreamPackLowBitValues8Bit( - attr.getRawData(), apply); - } - } else { - if (is_4bit_data) { - return tflite::StreamPackLowBitValues( - attr.getValues(), apply); - } else { - return tflite::StreamPackLowBitValues( - attr.getValues(), apply); - } - } + return PackLowBitElementsAttr(attr_and_inst.first, low_bit_width, + apply); }; } else { applier = @@ -1204,7 +1274,9 @@ std::optional> Translator::BuildBuffer( // is big endian, rely on the TensorFlow path below to reverse the // byte order. if (llvm::sys::IsLittleEndianHost && shaped_type && - shaped_type.getElementType().isIntOrFloat()) { + (shaped_type.getElementType().isIntOrFloat() || + mlir::isa( + shaped_type.getElementType()))) { int64_t expected_size = mlir::TFL::GetSizeInBytes(shaped_type); // DenseElementsAttr @@ -1222,13 +1294,11 @@ std::optional> Translator::BuildBuffer( // DenseResourceElementsAttr if (auto res_attr = mlir::dyn_cast(attr)) { - if (auto blob = - res_attr.getRawHandle().getResource()->getBlob()) { - auto data = blob->getData(); - if (data.size() == expected_size) { - return apply(absl::string_view( - reinterpret_cast(data.data()), data.size())); - } + mlir::AsmResourceBlob* blob = GetBlob(res_attr); + if (blob && blob->getData().size() == expected_size) { + return apply(absl::string_view( + reinterpret_cast(blob->getData().data()), + blob->getData().size())); } } } @@ -1291,9 +1361,10 @@ std::optional> Translator::BuildBuffer( // string and computing the hash of the string, but can be reliable in some // cases where the MLIR attributes are not deduped properly (e.g. when two // consts of the same value are held in different attribute types). + uint64_t h = GetPhysicalBufferHash(attr); const_buffer_storage_.Insert( index, std::make_pair(attr, inst), std::move(applier), - /*hash=*/mlir::hash_value(attr), + /*hash=*/h, /*byte_size_hint=*/mlir::TFL::GetSizeInBytes(type)); return tflite::CreateBuffer(builder_, 0, 1, 1); } else { @@ -3640,6 +3711,12 @@ std::optional> Translator::BuildSubGraph( for (auto result : bb.getTerminator()->getOperands()) { outputs.push_back(tensor_index_map[result]); } + if (index >= subgraph_inputs_.size()) { + subgraph_inputs_.resize(index + 1); + subgraph_outputs_.resize(index + 1); + } + subgraph_inputs_[index] = inputs; + subgraph_outputs_[index] = outputs; for (const auto& [from, to] : control_edges) { for (int what : {from, to}) { if (operation_index_to_operator_index.count(what) == 0) { @@ -3987,7 +4064,9 @@ std::vector GetStringsFromDictionaryAttr( std::vector BuildSignaturedef( FuncOp main_op, const std::string& saved_model_tag, - const uint32_t subgraph_index, tensorflow::OpOrArgNameMapper& name_mapper) { + const uint32_t subgraph_index, + const std::vector& input_tensor_indices, + const std::vector& output_tensor_indices) { static const char kEntryFunctionAttributes[] = "tf.entry_function"; // Fetch inputs and outputs from the signature. @@ -4050,14 +4129,10 @@ std::vector BuildSignaturedef( // We create vector of size 1 as TFLite now supports only 1 signatureDef. std::vector result(1); for (int i = 0; i < input_names.size(); ++i) { - result[0].inputs[sig_def_inputs[i]] = input_names[i].str(); + result[0].inputs[sig_def_inputs[i]] = input_tensor_indices[i]; } for (int i = 0; i < output_names.size(); ++i) { - // Fetch the name from the actual operand and not rely on names from - // outputs as deduping can make them invalid after conversion. - auto& operand = term->getOpOperand(i); - auto unique_name = std::string(name_mapper.GetUniqueName(operand.get())); - result[0].outputs[sig_def_outputs[i]] = unique_name; + result[0].outputs[sig_def_outputs[i]] = output_tensor_indices[i]; } if (auto name_attr = mlir::dyn_cast_or_null(exported_name[0])) result[0].signature_key = name_attr.getValue().str(); @@ -4066,14 +4141,13 @@ std::vector BuildSignaturedef( } std::vector> Translator::GetList( - const int subgraph_index, const std::map& items) { + const std::map& items) { std::vector> result; for (const auto& item : items) { auto name_buf = builder_.CreateString(item.first); tflite::TensorMapBuilder tensor_map_builder(builder_); tensor_map_builder.add_name(name_buf); - tensor_map_builder.add_tensor_index( - tensor_index_map_[subgraph_index][item.second]); + tensor_map_builder.add_tensor_index(item.second); result.push_back(tensor_map_builder.Finish()); } return result; @@ -4083,14 +4157,9 @@ std::optional>> Translator::CreateSignatureDefs( const std::vector& signature_defs) { std::vector> signature_defs_buffer; - // When we export each function in the module op, intentionally, we export - // the entry functions at the beginning of the subgraph list and the - // subgraph_index is the index in entry functions and at the same, is the - // index in the subgraph list. - int subgraph_index = 0; for (const auto& signature_def_data : signature_defs) { - auto inputs = GetList(subgraph_index, signature_def_data.inputs); - auto outputs = GetList(subgraph_index, signature_def_data.outputs); + auto inputs = GetList(signature_def_data.inputs); + auto outputs = GetList(signature_def_data.outputs); auto inputs_buf = builder_.CreateVector(inputs); auto outputs_buf = builder_.CreateVector(outputs); auto signature_key_buf = @@ -4101,7 +4170,6 @@ Translator::CreateSignatureDefs( sig_def_builder.add_signature_key(signature_key_buf); sig_def_builder.add_subgraph_index(signature_def_data.subgraph_index); signature_defs_buffer.push_back(sig_def_builder.Finish()); - ++subgraph_index; } return builder_.CreateVector(signature_defs_buffer); @@ -4144,9 +4212,11 @@ absl::Status Translator::Translate( op_or_arg_name_mapper = &default_op_or_arg_name_mapper; } if (!UpdateEntryFunction(module)) { + LOG(ERROR) << "No entry function found in the module."; return absl::InvalidArgumentError("No entry function found."); } if (!IsValidTFLiteMlirModule(module)) { + LOG(ERROR) << "Invalid TFLite MLIR module."; return absl::InvalidArgumentError("Invalid TFLite MLIR module."); } @@ -4369,7 +4439,8 @@ absl::Status Translator::TranslateInternal() { for (auto fn : entry_functions) { auto signature_defs = BuildSignaturedef( fn, saved_model_tags_.empty() ? "" : *saved_model_tags_.begin(), - subgraph_index, name_mapper_); + subgraph_index, subgraph_inputs_[subgraph_index], + subgraph_outputs_[subgraph_index]); for (const auto& signature_def : signature_defs) { signature_defs_vec.push_back(signature_def); } @@ -4721,6 +4792,7 @@ bool MlirToFlatBufferTranslateFunction(mlir::ModuleOp module, } if (!status.ok()) { + LOG(ERROR) << "Flatbuffer export failed: " << status.message(); return false; } serialized_flatbuffer->assign(buffer.data(), buffer.size()); diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc index f4fd0b5f13a409..61230aa335c82c 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc @@ -260,7 +260,9 @@ bool HasDenseResourceOperand(mlir::Operation* op) { // TODO(b/394905516): Remove this once we have a way to configure the threshold. bool ShouldFoldOperation(Operation* inst) { if (!(ENABLE_DENSE_RESOURCE_ATTR_FOLD) && HasDenseResourceOperand(inst)) { - return false; + if (!llvm::isa(inst)) { + return false; + } } auto get_size = [&](TypeRange types) { diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td index 4f03db9a996976..b281b8bee7748d 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td @@ -4397,7 +4397,7 @@ def TFL_QConstOp : Op:$output); + let results = (outs TFL_TensorOf<[QUI8, QI8, QUI4, QI4, QI16, QUI16, TFL_Quint8]>:$output); let builders = [ OpBuilder<(ins "TypeAttr":$qtype, "Attribute":$value), @@ -4428,7 +4428,7 @@ def TFL_SparseQConstOp : Op:$output); + let results = (outs TFL_TensorOf<[QUI8, QI8, QUI4, QI4, QI16, QUI16, TFL_Quint8]>:$output); let builders = [ OpBuilder<(ins "TypeAttr":$qtype, "Attribute":$value, diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/prepare_hlo.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/prepare_hlo.mlir index 6b6703285cec7f..36d1e91d8b0f26 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/tests/prepare_hlo.mlir +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/prepare_hlo.mlir @@ -102,6 +102,27 @@ func.func @depthwise_transpose_conv2d_same_padding_nchw_hwoi(%input: tensor<1x2x // ----- +// CHECK-LABEL: grouped_transpose_conv2d_nhwc_ohwi +func.func @grouped_transpose_conv2d_nhwc_ohwi(%input: tensor<1x8x8x4xf32>, %filter: tensor<8x3x3x2xf32>) -> tensor<1x17x17x8xf32> { + %0 = mhlo.convolution(%input, %filter) + dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f], + window = {pad = [[2, 2], [2, 2]], lhs_dilate = [2, 2]} + {batch_group_count = 1 : i64, feature_group_count = 2 : i64} + : (tensor<1x8x8x4xf32>, tensor<8x3x3x2xf32>) -> tensor<1x17x17x8xf32> + func.return %0 : tensor<1x17x17x8xf32> + + // CHECK: %0 = "mhlo.slice"(%arg0) <{limit_indices = dense<[1, 8, 8, 2]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<1x8x8x4xf32>) -> tensor<1x8x8x2xf32> + // CHECK: %1 = "mhlo.slice"(%arg1) <{limit_indices = dense<[4, 3, 3, 2]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<8x3x3x2xf32>) -> tensor<4x3x3x2xf32> + // CHECK: %2 = mhlo.convolution(%0, %1) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f], window = {pad = {{\[\[}}2, 2], [2, 2]], lhs_dilate = [2, 2]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x8x2xf32>, tensor<4x3x3x2xf32>) -> tensor<1x17x17x4xf32> + // CHECK: %3 = "mhlo.slice"(%arg0) <{limit_indices = dense<[1, 8, 8, 4]> : tensor<4xi64>, start_indices = dense<[0, 0, 0, 2]> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<1x8x8x4xf32>) -> tensor<1x8x8x2xf32> + // CHECK: %4 = "mhlo.slice"(%arg1) <{limit_indices = dense<[8, 3, 3, 2]> : tensor<4xi64>, start_indices = dense<[4, 0, 0, 0]> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<8x3x3x2xf32>) -> tensor<4x3x3x2xf32> + // CHECK: %5 = mhlo.convolution(%3, %4) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f], window = {pad = {{\[\[}}2, 2], [2, 2]], lhs_dilate = [2, 2]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x8x2xf32>, tensor<4x3x3x2xf32>) -> tensor<1x17x17x4xf32> + // CHECK: %6 = "mhlo.concatenate"(%2, %5) <{dimension = 3 : i64}> : (tensor<1x17x17x4xf32>, tensor<1x17x17x4xf32>) -> tensor<1x17x17x8xf32> + // CHECK: return %6 : tensor<1x17x17x8xf32> +} + +// ----- + // CHECK-LABEL: conv2d_nhwc_ohwi_nhwc func.func @conv2d_nhwc_ohwi_nhwc(%input: tensor<1x256x256x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32> { %0 = mhlo.convolution(%input, %filter) diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/conv.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/conv.cc index 11271daaec5654..c20db8b758a943 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/conv.cc +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/conv.cc @@ -46,25 +46,6 @@ using ::llvm::ArrayRef; // support/legality checking //===----------------------------------------------------------------------===// -bool IsShapeFullyStatic(ArrayRef shape) { - return llvm::all_of(shape, [](int64_t d) { return d >= 0; }); -} - -bool NonBatchDimsFullyStatic(ArrayRef shape) { - return IsShapeFullyStatic(shape.drop_front()); -} - -bool AreShapesFullyStatic(const ConvView& data) { - return IsShapeFullyStatic(data.InputShape()) && - IsShapeFullyStatic(data.KernelShape()) && - IsShapeFullyStatic(data.OutputShape()); -} - -bool InputOutputNonBatchDimsFullyStatic(const ConvView& data) { - return NonBatchDimsFullyStatic(data.InputShape()) && - IsShapeFullyStatic(data.KernelShape()) && - NonBatchDimsFullyStatic(data.OutputShape()); -} bool IsPaddingSupported(const ConvView& data) { return llvm::all_of(data.Padding(), [](const DimPadding& p) { @@ -453,7 +434,7 @@ LogicalResult ConvertNonTrivialConvToTransposeConvOp::matchAndRewrite( //===----------------------------------------------------------------------===// -class SliceDepthwiseTransposedConvolution +class SliceGroupedTransposedConvolution : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; @@ -461,12 +442,12 @@ class SliceDepthwiseTransposedConvolution PatternRewriter& rewriter) const final; }; -// Pattern rewriter to match a depthwise transposed convolution and rewrite it -// to depth-times slices of input and filter to perform the transposed -// convolution on individual slices of tensors and concatenate the results of. -// the convolutions. This is a. workaround because the TFLite runtime doesn't -// support depthwise-transposed-conv op natively. -LogicalResult SliceDepthwiseTransposedConvolution::matchAndRewrite( +// Pattern rewriter to match a grouped (or depthwise) transposed convolution and +// rewrite it to group-times slices of input and filter to perform the +// transposed convolution on individual slices of tensors and concatenate the +// results of the convolutions. This is a workaround because the TFLite runtime +// doesn't support grouped or depthwise-transposed-conv op natively. +LogicalResult SliceGroupedTransposedConvolution::matchAndRewrite( mhlo::ConvolutionOp conv_op, PatternRewriter& rewriter) const { const ConvView data(conv_op); @@ -478,7 +459,13 @@ LogicalResult SliceDepthwiseTransposedConvolution::matchAndRewrite( "Not a non-trivial convolution."); } - // These checks narrow down the support to depthwise transpose conv2d. + if (!mlir::cast(conv_op.getLhs().getType()).hasStaticShape() || + !mlir::cast(conv_op.getRhs().getType()).hasStaticShape() || + !mlir::cast(conv_op.getType()).hasStaticShape()) { + return rewriter.notifyMatchFailure(conv_op, "Requires static shapes."); + } + + // These checks narrow down the support to grouped transpose conv2d. mhlo::ConvDimensionNumbersAttr dnums = conv_op.getDimensionNumbers(); const int64_t input_feature_dimension = dnums.getInputFeatureDimension(); const int64_t input_channels = @@ -496,60 +483,49 @@ LogicalResult SliceDepthwiseTransposedConvolution::matchAndRewrite( mlir::cast(conv_op.getRhs().getType()) .getDimSize(kernel_output_feature_dimension); - // To support a depthwise convolution, we need- - // 1. feature_group_count != 1 (except when input_channels==1) - // 2. feature_group_count == input_channels - // 3. kernel_input_channels == 1 - // 4. kernel_output_channels % kernel_input_channels == 0 - if (feature_group_count == 1) { - return rewriter.notifyMatchFailure(conv_op, "Not a depthwise convolution"); + if (feature_group_count <= 1) { + return rewriter.notifyMatchFailure(conv_op, + "Not a grouped transposed convolution"); } - if (input_channels != feature_group_count) { + if (input_channels % feature_group_count != 0) { return rewriter.notifyMatchFailure( - conv_op, "Not a detphwise transposed convolution"); + conv_op, "Input channels not divisible by feature group count"); } - if (MatchWithResizeBilinearOp(data)) { + const int64_t in_channels_per_group = input_channels / feature_group_count; + if (kernel_input_channels != in_channels_per_group) { return rewriter.notifyMatchFailure( - conv_op, "Op will be legalized to ResizeBilinearOp"); + conv_op, + "Kernel input channels does not match input channels per group"); } - if ((kernel_output_channels % feature_group_count != 0) || - (kernel_input_channels != 1)) { + if (kernel_output_channels % feature_group_count != 0) { return rewriter.notifyMatchFailure( - conv_op, "Not a supported detphwise transposed convolution"); + conv_op, "Kernel output channels not divisible by feature group count"); } - // This needs to be checked because the TFLite runtime generated incorrect - // results for depthwise transpose convolutions with non-1 channel - // multiplier. - if ((kernel_output_channels / feature_group_count) != 1) { + const int64_t out_channels_per_group = + kernel_output_channels / feature_group_count; + + if (MatchWithResizeBilinearOp(data)) { return rewriter.notifyMatchFailure( - conv_op, - "Unsupported detphwise transpose convolution with non-1 channel " - "multiplier"); + conv_op, "Op will be legalized to ResizeBilinearOp"); } // Slicing with dynamic offsets (helper method advised) - auto create_slice = [&](mlir::Value tensor, int64_t depth_idx, + auto create_slice = [&](mlir::Value tensor, int64_t group_idx, int64_t channel_idx, - bool is_kernel = false) -> mlir::Value { + int64_t channels_per_group) -> mlir::Value { auto tensor_shape = mlir::cast(tensor.getType()).getShape().vec(); - // Calculate offsets based on depth_idx, channel_idx and tensor_shape + // Calculate offsets based on group_idx, channel_idx and tensor_shape llvm::SmallVector start_indices(tensor_shape.size(), 0); auto limit_indices = tensor_shape; const llvm::SmallVector strides(tensor_shape.size(), 1); - start_indices[channel_idx] = depth_idx; - if (is_kernel) { - // kernel can have a channel_multiplier that needs to be accounted for - limit_indices[channel_idx] = - depth_idx + (kernel_output_channels / feature_group_count); - } else { - limit_indices[channel_idx] = depth_idx + 1; - } + start_indices[channel_idx] = group_idx * channels_per_group; + limit_indices[channel_idx] = (group_idx + 1) * channels_per_group; return mhlo::SliceOp::create(rewriter, conv_op.getLoc(), tensor, rewriter.getI64TensorAttr(start_indices), rewriter.getI64TensorAttr(limit_indices), @@ -561,16 +537,18 @@ LogicalResult SliceDepthwiseTransposedConvolution::matchAndRewrite( // Iterative Slicing and Convolutions for (int i = 0; i < feature_group_count; ++i) { - auto sliced_input = - create_slice(conv_op.getLhs(), i, input_feature_dimension); - auto sliced_kernel = create_slice(conv_op.getRhs(), i, - kernel_output_feature_dimension, true); + auto sliced_input = create_slice( + conv_op.getLhs(), i, input_feature_dimension, in_channels_per_group); + auto sliced_kernel = + create_slice(conv_op.getRhs(), i, kernel_output_feature_dimension, + out_channels_per_group); // Calculate convolution output_type based on sliced_input and // sliced_kernel auto output_type = mlir::cast(conv_op->getResult(0).getType()); auto new_output_shape = output_type.getShape().vec(); - new_output_shape[dnums.getOutputFeatureDimension()] /= feature_group_count; + new_output_shape[dnums.getOutputFeatureDimension()] = + out_channels_per_group; auto new_output_type = RankedTensorType::get(new_output_shape, output_type.getElementType()); @@ -766,6 +744,6 @@ void PopulateLegalizeConvPatterns(MLIRContext* ctx, RewritePatternSet& patterns, void PopulatePrepareConvPatterns(MLIRContext* ctx, RewritePatternSet& patterns) { - patterns.add(ctx); + patterns.add(ctx); } } // namespace mlir::odml diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dense_constants.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dense_constants.mlir index d8b4e6b7b4fa75..f280182eb0afae 100644 --- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dense_constants.mlir +++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dense_constants.mlir @@ -51,6 +51,13 @@ func.func @uint8() -> tensor<4xui8> { func.return %0 : tensor<4xui8> } +func.func @empty() -> tensor<0xi32> { + // CHECK-LABEL: @empty + // CHECK: value = dense<> : tensor<0xi32> + %0 = "tfl.pseudo_const"() { value = dense_resource : tensor<0xi32> } : () -> tensor<0xi32> + func.return %0 : tensor<0xi32> +} + // Identity function to make the exporter happy func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> { func.return %arg0 : tensor<4xi8> @@ -63,7 +70,8 @@ func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> { dense_elements_i16: "0x400000000100020003000201", dense_elements_i32: "0x4000000001000000020000000300000004030201", dense_elements_i8: "0x4000000001020304", - dense_elements_i8_1: "0x40000000DEADBEEF" + dense_elements_i8_1: "0x40000000DEADBEEF", + dense_elements_empty: "0x40000000" } } #-} diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/low_bit_packing.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/low_bit_packing.mlir index 8b6b3bffd1ccab..c65b909b1cc252 100644 --- a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/low_bit_packing.mlir +++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/low_bit_packing.mlir @@ -14,7 +14,7 @@ // ============================================================================== // RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -emit-custom-ops -emit-builtin-tflite-ops=false -o - | flatbuffer_to_string - | FileCheck %s -func.func @main() -> tensor<4xi4> { +func.func @main() -> (tensor<4xi4>, tensor<4xi4>) { // CHECK: { // CHECK: version: 3, // CHECK: operator_codes: [ ], @@ -28,9 +28,18 @@ func.func @main() -> tensor<4xi4> { // CHECK-EMPTY // CHECK: }, // CHECK: has_rank: true + // CHECK: }, { + // CHECK: shape: [ 4 ], + // CHECK: type: INT4, + // CHECK: buffer: 2, + // CHECK: name: "ConstResource", + // CHECK: quantization: { + // CHECK-EMPTY + // CHECK: }, + // CHECK: has_rank: true // CHECK: } ], // CHECK: inputs: [ ], - // CHECK: outputs: [ 0 ], + // CHECK: outputs: [ 0, 1 ], // CHECK: operators: [ ], // CHECK: name: "main" // CHECK: } ], @@ -40,11 +49,13 @@ func.func @main() -> tensor<4xi4> { // CHECK: }, { // CHECK: data: [ 56, 190 ] // CHECK: }, { + // CHECK: data: [ 56, 190 ] + // CHECK: }, { // CHECK: data: [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] // CHECK: } ], // CHECK: metadata: [ { // CHECK: name: "min_runtime_version", - // CHECK: buffer: 2 + // CHECK: buffer: 3 // CHECK: } ], // CHECK: signature_defs: [ ] // CHECK: } @@ -53,5 +64,14 @@ func.func @main() -> tensor<4xi4> { // be packed low-bits-first as [0x38, 0xBE] or [56, 190]. Tensor type should // be INT4. %0 = "tfl.pseudo_const" () {value = dense<[-8, 3, -2, -5]> : tensor<4xi4>} : () -> tensor<4xi4> loc("Const") - func.return %0 : tensor<4xi4> + %1 = "tfl.pseudo_const" () {value = dense_resource : tensor<4xi4>} : () -> tensor<4xi4> loc("ConstResource") + func.return %0, %1 : tensor<4xi4>, tensor<4xi4> } + +{-# + dialect_resources: { + builtin: { + res_i4: "0x40000000F803FEFB" + } + } +#-} diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/signature_def_same_name.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/signature_def_same_name.mlir new file mode 100644 index 00000000000000..75a655fea0e9dc --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/signature_def_same_name.mlir @@ -0,0 +1,86 @@ +// Copyright 2026 The TensorFlow Authors. All Rights Reserved. +// +// 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. +// ============================================================================== +// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s + +// Verify that when an input and output share the same name ("x"), +// the signature_def maps the input to tensor index 0 and the output to +// the computed output tensor index (1) rather than aliasing to 0. + +// CHECK: { +// CHECK-NEXT: version: 3, +// CHECK-NEXT: operator_codes: [ { +// CHECK-NEXT: version: 1 +// CHECK-NEXT: } ], +// CHECK-NEXT: subgraphs: [ { +// CHECK-NEXT: tensors: [ { +// CHECK-NEXT: shape: [ 1, 4 ], +// CHECK-NEXT: buffer: 1, +// CHECK-NEXT: name: "x", +// CHECK-NEXT: quantization: { +// CHECK-EMPTY: +// CHECK-NEXT: }, +// CHECK-NEXT: has_rank: true +// CHECK-NEXT: }, { +// CHECK-NEXT: shape: [ 1, 4 ], +// CHECK-NEXT: buffer: 2, +// CHECK-NEXT: name: "x", +// CHECK-NEXT: quantization: { +// CHECK-EMPTY: +// CHECK-NEXT: }, +// CHECK-NEXT: has_rank: true +// CHECK-NEXT: } ], +// CHECK-NEXT: inputs: [ 0 ], +// CHECK-NEXT: outputs: [ 1 ], +// CHECK-NEXT: operators: [ { +// CHECK-NEXT: inputs: [ 0, 0 ], +// CHECK-NEXT: outputs: [ 1 ], +// CHECK-NEXT: builtin_options_type: AddOptions, +// CHECK-NEXT: builtin_options: { +// CHECK-EMPTY: +// CHECK-NEXT: } +// CHECK-NEXT: } ], +// CHECK-NEXT: name: "main" +// CHECK-NEXT: } ], +// CHECK-NEXT: description: "MLIR Converted.", +// CHECK-NEXT: buffers: [ { +// CHECK-EMPTY: +// CHECK-NEXT: }, { +// CHECK-EMPTY: +// CHECK-NEXT: }, { +// CHECK-EMPTY: +// CHECK-NEXT: }, { +// CHECK-NEXT: data: [ 49, 46, 53, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] +// CHECK-NEXT: } ], +// CHECK-NEXT: metadata: [ { +// CHECK-NEXT: name: "min_runtime_version", +// CHECK-NEXT: buffer: 3 +// CHECK-NEXT: } ], +// CHECK-NEXT: signature_defs: [ { +// CHECK-NEXT: inputs: [ { +// CHECK-NEXT: name: "x" +// CHECK-NEXT: } ], +// CHECK-NEXT: outputs: [ { +// CHECK-NEXT: name: "x", +// CHECK-NEXT: tensor_index: 1 +// CHECK-NEXT: } ], +// CHECK-NEXT: signature_key: "serving_default" +// CHECK-NEXT: } ] +// CHECK-NEXT:} +module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} { + func.func @main(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "x", outputs = "x"}, tf_saved_model.exported_names = ["serving_default"]} { + %0 = "tfl.add"(%arg0, %arg0) {fused_activation_function = "NONE"} : (tensor<1x4xf32>, tensor<1x4xf32>) -> tensor<1x4xf32> + func.return %0 : tensor<1x4xf32> + } +} diff --git a/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.cc b/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.cc index 4ae658ba39e21e..69f4d8594cf99c 100644 --- a/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.cc +++ b/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.cc @@ -489,40 +489,32 @@ int64_t GetSizeInBits(mlir::ShapedType shaped_type) { } int64_t GetSizeInBits(mlir::quant::QuantizedType quant_type) { - const int64_t bits = std::max(quant_type.getStorageTypeIntegralWidth(), - static_cast(CHAR_BIT)); + const int64_t bits = quant_type.getStorageTypeIntegralWidth(); assert(IsPowerOfTwo(bits)); return bits; } int64_t GetSizeInBits(mlir::Type type) { if (type.isIntOrFloat()) { - const int64_t bits = - std::max(type.getIntOrFloatBitWidth(), static_cast(CHAR_BIT)); + const int64_t bits = type.getIntOrFloatBitWidth(); assert(IsPowerOfTwo(bits)); return bits; } - if (mlir::isa(type)) { - auto shaped_type = mlir::cast(type); - if (mlir::isa(shaped_type.getElementType())) { - auto complex_type = - mlir::cast(shaped_type.getElementType()); - return GetSizeInBits(complex_type.getElementType()) * 2; - } else if (mlir::isa( - shaped_type.getElementType())) { - auto quant_type = - mlir::cast(shaped_type.getElementType()); - return GetSizeInBits(quant_type); - } else { - return GetSizeInBits(shaped_type); - } + if (auto complex_type = mlir::dyn_cast(type)) { + return GetSizeInBits(complex_type.getElementType()) * 2; + } + if (auto quant_type = mlir::dyn_cast(type)) { + return GetSizeInBits(quant_type); + } + if (auto shaped_type = mlir::dyn_cast(type)) { + return GetSizeInBits(shaped_type); } return 0; } int64_t GetSizeInBytes(mlir::Type type) { - return ExactIntegerDivide(GetSizeInBits(type), CHAR_BIT); + return (GetSizeInBits(type) + CHAR_BIT - 1) / CHAR_BIT; } } // namespace TFL diff --git a/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h b/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h index 7037fa2d322d76..414941dbc1e45c 100644 --- a/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h +++ b/tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h @@ -72,16 +72,12 @@ tensorflow::TensorProto ConvertTfliteConstTensor( // Get the size of the type in bits. The type can be ComplexType, FloatType, // IntegerType, QuantizedType, or ShapeType of other supported types. -// -// Sub-byte types, e.g. qu4 and i2, are treated as a full i8. int64_t GetSizeInBits(mlir::ShapedType shaped_type); int64_t GetSizeInBits(mlir::Type type); int64_t GetSizeInBits(mlir::quant::QuantizedType quant_type); -// Get the size of the type in bytes. -// -// Sub-byte element types, e.g. qu4 and i2, are treated as a full i8. -// e.g. GetSizeInBytes(tensor<4xi2>) == 4, instead of 1. +// Get the size of the type in bytes (rounded up for sub-byte types). +// e.g. GetSizeInBytes(tensor<4xi2>) == 1. int64_t GetSizeInBytes(mlir::Type type); // Performs an integer divide and checks that the remainder is zero. diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index f111426835174d..92786fc8f30ed7 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -485,7 +485,6 @@ cc_library( "@com_google_absl//absl/types:span", ":encoded_buffer_allocation_info", "@xla//xla/service:custom_call_status_internal", - "@xla//xla/backends/cpu/runtime:msan_emulated_tls", "@xla//xla/backends/cpu/runtime:rng_state_lib", "@xla//xla/backends/cpu:alignment", "@xla//xla/backends/cpu:buffer_allocation_info", diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD index 64c4ce9b8cb160..31766e878ec9e9 100644 --- a/tensorflow/core/common_runtime/BUILD +++ b/tensorflow/core/common_runtime/BUILD @@ -3023,6 +3023,7 @@ tf_cc_test( srcs = ["process_util_test.cc"], deps = [ ":process_util", + "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", ], diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc index e3df3c15b941b1..ef72a33d1ffa04 100644 --- a/tensorflow/core/common_runtime/direct_session.cc +++ b/tensorflow/core/common_runtime/direct_session.cc @@ -135,7 +135,7 @@ absl::Status NewThreadPoolFromThreadPoolOptions( const ThreadPoolOptionProto& thread_pool_options, int pool_number, thread::ThreadPool** pool, bool* owned) { int32_t num_threads = thread_pool_options.num_threads(); - if (num_threads == 0) { + if (num_threads <= 0) { num_threads = NumInterOpThreadsFromSessionOptions(options); } const std::string& name = thread_pool_options.global_name(); diff --git a/tensorflow/core/common_runtime/local_device.cc b/tensorflow/core/common_runtime/local_device.cc index 9997ff2a30c008..64972e4c076654 100644 --- a/tensorflow/core/common_runtime/local_device.cc +++ b/tensorflow/core/common_runtime/local_device.cc @@ -73,11 +73,11 @@ struct LocalDevice::EigenThreadPoolInfo { int32_t intra_op_parallelism_threads = options.config.intra_op_parallelism_threads(); // If no session setting, use environment setting. - if (intra_op_parallelism_threads == 0) { + if (intra_op_parallelism_threads <= 0) { static int env_num_threads = NumIntraOpThreadsFromEnvironment(); intra_op_parallelism_threads = env_num_threads; // If no session setting or environment, compute a reasonable default. - if (intra_op_parallelism_threads == 0) { + if (intra_op_parallelism_threads <= 0) { intra_op_parallelism_threads = port::MaxParallelism(numa_node); } } diff --git a/tensorflow/core/common_runtime/process_util.cc b/tensorflow/core/common_runtime/process_util.cc index 233dcde498a6bc..4ec5548627aa99 100644 --- a/tensorflow/core/common_runtime/process_util.cc +++ b/tensorflow/core/common_runtime/process_util.cc @@ -74,7 +74,7 @@ int32_t DefaultNumInterOpThreads() { static thread::ThreadPool* InitComputePool(const SessionOptions& options) { int32_t inter_op_parallelism_threads = options.config.inter_op_parallelism_threads(); - if (inter_op_parallelism_threads == 0) { + if (inter_op_parallelism_threads <= 0) { inter_op_parallelism_threads = DefaultNumInterOpThreads(); } return new thread::ThreadPool( diff --git a/tensorflow/core/common_runtime/process_util_test.cc b/tensorflow/core/common_runtime/process_util_test.cc index 46672ac92eef1c..d1a53e8f411854 100644 --- a/tensorflow/core/common_runtime/process_util_test.cc +++ b/tensorflow/core/common_runtime/process_util_test.cc @@ -14,6 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/common_runtime/process_util.h" +#include + +#include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { @@ -29,9 +32,18 @@ TEST(ProcessUtilTest, ThreadPool) { SessionOptions opts; opts.config.set_inter_op_parallelism_threads(10); - thread::ThreadPool* pool = NewThreadPoolFromSessionOptions(opts); + std::unique_ptr pool( + NewThreadPoolFromSessionOptions(opts)); EXPECT_EQ(10, pool->NumThreads()); - delete pool; +} + +TEST(ProcessUtilTest, ThreadPoolNegative) { + SessionOptions opts; + opts.config.set_inter_op_parallelism_threads(-1); + + std::unique_ptr pool( + NewThreadPoolFromSessionOptions(opts)); + EXPECT_GT(pool->NumThreads(), 0); } } // anonymous namespace diff --git a/tensorflow/core/kernels/linalg/determinant_op.cc b/tensorflow/core/kernels/linalg/determinant_op.cc index ed9fd5f061a2bb..51db773fa289a5 100644 --- a/tensorflow/core/kernels/linalg/determinant_op.cc +++ b/tensorflow/core/kernels/linalg/determinant_op.cc @@ -377,7 +377,7 @@ class LogDeterminantOpGpu : public AsyncOpKernel { // input_copy by the Getrf{Batched} kernel. functor::LogDeterminantFromPivotedLUFunctor functor; functor(d, input_copy_reshaped_const, pivots_mat.data(), sign_reshaped, - log_abs_det_reshaped); + log_abs_det_reshaped, dev_info.back().mutable_data()); // Register callback to check info after kernels finish. auto info_checker = [context, done]( @@ -399,6 +399,7 @@ class LogDeterminantOpGpu : public AsyncOpKernel { } done(); }; + GpuSolver::CheckLapackInfoAndDeleteSolverAsync(std::move(solver), dev_info, std::move(info_checker)); } diff --git a/tensorflow/core/kernels/linalg/determinant_op.h b/tensorflow/core/kernels/linalg/determinant_op.h index 6ace1bef44b250..47de2dbccf010b 100644 --- a/tensorflow/core/kernels/linalg/determinant_op.h +++ b/tensorflow/core/kernels/linalg/determinant_op.h @@ -28,7 +28,7 @@ struct DeterminantFromPivotedLUFunctor { void operator()(const Device& device, typename TTypes::ConstTensor lu_factor, const int* pivots, typename TTypes::Tensor output, - int* info); + const int* info); }; // Helper functor to compute sign and log of the absolute value of the @@ -38,7 +38,8 @@ struct LogDeterminantFromPivotedLUFunctor { void operator()(const Device& device, typename TTypes::ConstTensor lu_factor, const int* pivots, typename TTypes::Tensor sign, - typename TTypes::Tensor log_abs_det); + typename TTypes::Tensor log_abs_det, + const int* info); }; } // namespace functor diff --git a/tensorflow/core/kernels/linalg/determinant_op_gpu.cu.cc b/tensorflow/core/kernels/linalg/determinant_op_gpu.cu.cc index 5dff239561dfde..69048032149a37 100644 --- a/tensorflow/core/kernels/linalg/determinant_op_gpu.cu.cc +++ b/tensorflow/core/kernels/linalg/determinant_op_gpu.cu.cc @@ -49,8 +49,8 @@ __device__ int PermutationOrder(int n, const int* __restrict__ pivots) { template __global__ void DeterminantFromPivotedLUKernel( int nthreads, int n, const Scalar* __restrict__ lu_factor, - const int* __restrict__ all_pivots, Scalar* __restrict__ sign, - Scalar* __restrict__ log_abs_det) { + const int* __restrict__ all_pivots, const int* __restrict__ info, + Scalar* __restrict__ sign, Scalar* __restrict__ log_abs_det) { typedef typename Eigen::NumTraits::Real RealScalar; const int matrix_size = n * n; const int stride = n + 1; @@ -59,6 +59,16 @@ __global__ void DeterminantFromPivotedLUKernel( // The main purpose is to avoid having to copy the LU decomposition to // host memory. GPU_1D_KERNEL_LOOP(o_idx, nthreads) { + if (info != nullptr && info[o_idx] > 0) { + if (compute_log_abs_det) { + sign[o_idx] = Scalar(0); + log_abs_det[o_idx] = + Scalar(-Eigen::numext::numeric_limits::infinity()); + } else { + log_abs_det[o_idx] = Scalar(0); + } + continue; + } // Initialize sign to (-1)^order. const int order = PermutationOrder(n, all_pivots + o_idx * n); Scalar prod_sign = order % 2 ? Scalar(-1) : Scalar(1); @@ -99,7 +109,7 @@ struct DeterminantFromPivotedLUFunctor { void operator()(const GPUDevice& device, typename TTypes::ConstTensor lu_factor, const int* pivots, typename TTypes::Tensor output, - int* info) { + const int* info) { const int64_t num_matrices = output.size(); const int64_t n = lu_factor.dimension(2); GpuLaunchConfig config = GetGpuLaunchConfig(num_matrices, device); @@ -107,8 +117,8 @@ struct DeterminantFromPivotedLUFunctor { TF_CHECK_OK(GpuLaunchKernel( DeterminantFromPivotedLUKernel, config.block_count, config.thread_per_block, 0, device.stream(), - config.virtual_thread_count, n, lu_factor.data(), pivots, nullptr, - output.data())); + config.virtual_thread_count, n, lu_factor.data(), pivots, info, + /*sign=*/nullptr, output.data())); } }; @@ -122,15 +132,16 @@ struct LogDeterminantFromPivotedLUFunctor { void operator()(const GPUDevice& device, typename TTypes::ConstTensor lu_factor, const int* pivots, typename TTypes::Tensor sign, - typename TTypes::Tensor log_abs_det) { + typename TTypes::Tensor log_abs_det, + const int* info) { const int64_t num_matrices = sign.size(); const int64_t n = lu_factor.dimension(2); GpuLaunchConfig config = GetGpuLaunchConfig(num_matrices, device); TF_CHECK_OK(GpuLaunchKernel( DeterminantFromPivotedLUKernel, config.block_count, config.thread_per_block, 0, device.stream(), - config.virtual_thread_count, n, lu_factor.data(), pivots, sign.data(), - log_abs_det.data())); + config.virtual_thread_count, n, lu_factor.data(), pivots, info, + sign.data(), log_abs_det.data())); } }; diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index ad7d9f8e4916b4..e358f4521598b1 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -1344,6 +1344,9 @@ class SparseSegmentGradOpBase : public OpKernel { OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); if (M == 0 || N == 0) return; + OP_REQUIRES(context, input.dim_size(0) > 0, + absl::InvalidArgumentError("Invalid number of segments")); + functor::SparseSegmentGradFunctor()( context, operation_, input_flat, indices_vec, segment_vec, output); } @@ -1422,6 +1425,10 @@ class SparseSegmentGradV2OpCommon { return absl::OkStatus(); } + if (input.dim_size(0) == 0) { + return absl::InvalidArgumentError("Invalid number of segments"); + } + auto input_flat = input.flat_outer_dims(); const auto indices_vec = indices.vec(); const auto segment_vec = segment_ids.vec(); diff --git a/tensorflow/lite/fuzzing/BUILD b/tensorflow/lite/fuzzing/BUILD deleted file mode 100644 index 223e7a35c0898b..00000000000000 --- a/tensorflow/lite/fuzzing/BUILD +++ /dev/null @@ -1,22 +0,0 @@ -# Fuzzing harnesses for the TensorFlow Lite runtime. - -load( - "//tensorflow/security/fuzzing:tf_fuzzing.bzl", - "tf_cc_fuzz_test", -) - -package( - # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"], - default_visibility = ["//visibility:private"], - licenses = ["notice"], -) - -tf_cc_fuzz_test( - name = "interpreter_fuzz", - srcs = ["interpreter_fuzz.cc"], - deps = [ - "//tensorflow/lite:framework", - "//tensorflow/lite/core:framework", - "//tensorflow/lite/kernels:builtin_ops", - ], -) diff --git a/tensorflow/lite/fuzzing/interpreter_fuzz.cc b/tensorflow/lite/fuzzing/interpreter_fuzz.cc deleted file mode 100644 index 785090ace761bd..00000000000000 --- a/tensorflow/lite/fuzzing/interpreter_fuzz.cc +++ /dev/null @@ -1,99 +0,0 @@ -/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. - -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. -==============================================================================*/ - -// Fuzzes the TensorFlow Lite interpreter end to end: an arbitrary buffer is -// verified as a flatbuffer model, an interpreter is built for it, tensors are -// allocated, inputs are filled deterministically, and the graph is invoked. -// -// This exercises the builtin kernel implementations, the arena planner and the -// shape-propagation paths, none of which previously had OSS-Fuzz coverage. - -#include -#include -#include -#include -#include - -#include "fuzztest/fuzztest.h" -#include "tensorflow/lite/core/interpreter.h" -#include "tensorflow/lite/core/interpreter_builder.h" -#include "tensorflow/lite/core/model_builder.h" -#include "tensorflow/lite/kernels/register.h" - -namespace tflite { -namespace fuzzing { -namespace { - -// Keep the fuzzer inside the OSS-Fuzz memory budget. Models and arenas larger -// than this are not interesting: they exercise the allocator, not the kernels. -constexpr size_t kMaxModelBytes = 1 << 20; // 1 MiB -constexpr size_t kMaxArenaBytes = 1 << 26; // 64 MiB - -void FuzzInterpreter(const std::string& model_bytes) { - if (model_bytes.size() < 8 || model_bytes.size() > kMaxModelBytes) { - return; - } - - // VerifyAndBuildFromBuffer applies tflite::VerifyModelBuffer first, so - // structurally invalid buffers are rejected cheaply. - std::unique_ptr model = - FlatBufferModel::VerifyAndBuildFromBuffer(model_bytes.data(), - model_bytes.size()); - if (model == nullptr) { - return; - } - - // Delegates are excluded so the fuzzer exercises the reference and optimized - // CPU kernels rather than a delegate's own implementation. - ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver; - std::unique_ptr interpreter; - if (InterpreterBuilder(*model, resolver)(&interpreter) != kTfLiteOk || - interpreter == nullptr) { - return; - } - - if (interpreter->AllocateTensors() != kTfLiteOk) { - return; - } - - size_t total_bytes = 0; - for (const int tensor_index : interpreter->inputs()) { - TfLiteTensor* tensor = interpreter->tensor(tensor_index); - if (tensor == nullptr || tensor->data.raw == nullptr) { - continue; - } - // String tensors own a dynamic buffer with its own layout; writing raw - // bytes into it would corrupt the interpreter rather than the kernel under - // test. - if (tensor->type == kTfLiteString || tensor->type == kTfLiteResource || - tensor->type == kTfLiteVariant) { - return; - } - // Check before accumulating so the sum itself cannot wrap. - if (tensor->bytes > kMaxArenaBytes || - total_bytes > kMaxArenaBytes - tensor->bytes) { - return; - } - total_bytes += tensor->bytes; - std::memset(tensor->data.raw, 1, tensor->bytes); - } - - interpreter->Invoke(); -} -FUZZ_TEST(TfLiteFuzz, FuzzInterpreter); - -} // namespace -} // namespace fuzzing -} // namespace tflite diff --git a/tensorflow/python/kernel_tests/linalg/determinant_op_test.py b/tensorflow/python/kernel_tests/linalg/determinant_op_test.py index e57947c34b1cfa..27f14bb59de69f 100644 --- a/tensorflow/python/kernel_tests/linalg/determinant_op_test.py +++ b/tensorflow/python/kernel_tests/linalg/determinant_op_test.py @@ -84,6 +84,27 @@ def testBasic(self): # A multidimensional batch of 2x2 matrices self._compareDeterminant(np.random.rand(3, 4, 5, 2, 2).astype(np.float32)) + def testSingularMatrix(self): + for dtype in (np.float32, np.float64, np.complex64, np.complex128): + for matrix_values in ( + [[1.0, 2.0], [2.0, 4.0]], + [[0.0, 1.0, 2.0], [0.0, 3.0, 4.0], [0.0, 5.0, 6.0]], + [[1.0, 0.0, -1.0], [-1.0, 1.0, 0.0], [0.0, -1.0, 1.0]], + [ + [[1.0, 2.0], [2.0, 4.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + ): + matrix = np.array(matrix_values, dtype=dtype) + with test_util.use_gpu(): + det = self.evaluate(linalg_ops.matrix_determinant(matrix)) + sign, log_det = self.evaluate( + gen_linalg_ops.log_matrix_determinant(matrix) + ) + self.assertAllClose(det, np.zeros_like(det)) + self.assertAllClose(sign, np.zeros_like(sign)) + self.assertTrue(np.all(np.isneginf(log_det.real))) + def testBasicDouble(self): # 2x2 matrices self._compareDeterminant(np.array([[2., 3.], [3., 4.]]).astype(np.float64)) diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 24cfb3eec109ec..2f7adecacfd6bf 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -1360,9 +1360,32 @@ def testGradientSegmentsInvalid4(self): with self.session(use_gpu=False): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) - with self.assertRaisesOpError(r"Segment id 0 out of range \[0, 0\)"): + with self.assertRaisesOpError("Invalid number of segments"): self.evaluate(s) + def testGradientEmptyInputWithNonEmptyIndices(self): + ops_list = [ + math_ops.sparse_segment_sum_grad, + math_ops.sparse_segment_mean_grad, + math_ops.sparse_segment_sqrt_n_grad, + ] + v2_ops_list = [ + math_ops.sparse_segment_sum_grad_v2, + math_ops.sparse_segment_mean_grad_v2, + math_ops.sparse_segment_sqrt_n_grad_v2, + ] + indices = [0, 1] + segment_ids = [0, 1] + output_dim0 = 2 + for dtype in [dtypes_lib.float16, dtypes_lib.float32, dtypes_lib.float64]: + grad = constant_op.constant([], shape=[0], dtype=dtype) + for tf_op in ops_list: + with self.assertRaisesOpError("Invalid number of segments"): + self.evaluate(tf_op(grad, indices, segment_ids, output_dim0)) + for tf_op in v2_ops_list: + with self.assertRaisesOpError("Invalid number of segments"): + self.evaluate(tf_op(grad, indices, segment_ids, output_dim0)) + def testGradientV2Valid(self): # Baseline for the testGradientV2*Invalid* methods below. tf_x, _ = self._input([3, 4], dtype=dtypes_lib.float32) @@ -1473,7 +1496,7 @@ def testGradientV2SegmentsInvalid4(self): with self.session(use_gpu=False): for tf_op in ops_list: s, j = tf_op(tf_x, tf_indices, segment_indices, 10) - with self.assertRaisesOpError(r"Segment id 0 out of range \[0, 0\)"): + with self.assertRaisesOpError("Invalid number of segments"): self.evaluate([s, j]) diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py index 7bfddf38185b5f..8ca1cd80c914f0 100644 --- a/tensorflow/python/training/server_lib_test.py +++ b/tensorflow/python/training/server_lib_test.py @@ -505,6 +505,16 @@ def testDenseAndSparseJobs(self): cluster_spec = server_lib.ClusterSpec(cluster_def) self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def()) + def testNegativeInterOpParallelismThreads(self): + config = config_pb2.ConfigProto(inter_op_parallelism_threads=-1) + server = server_lib.Server.create_local_server(config=config) + self.assertIsNotNone(server.target) + + def testNegativeIntraOpParallelismThreads(self): + config = config_pb2.ConfigProto(intra_op_parallelism_threads=-1) + server = server_lib.Server.create_local_server(config=config) + self.assertIsNotNone(server.target) + class ClusterSpecTest(test.TestCase): diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index dc2c4735acf69c..ebb9dda1528d7c 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -1,3 +1,20 @@ +# 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 file is automatically generated by generate_patch tool. +# Do not edit directly. diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md --- stablehlo/docs/spec.md +++ stablehlo/docs/spec.md diff --git a/third_party/xla/third_party/triton/oss_only/build_files.patch b/third_party/xla/third_party/triton/oss_only/build_files.patch index 7e6bba2cd06c6d..d98379f330bca5 100644 --- a/third_party/xla/third_party/triton/oss_only/build_files.patch +++ b/third_party/xla/third_party/triton/oss_only/build_files.patch @@ -1,3 +1,17 @@ +# 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 an autogenerated patch file. Do not edit directly. # To update, see instructions at go/patch-triton. diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD index 188a9299c67e8f..40f90b5b5e18df 100644 --- a/third_party/xla/xla/backends/cpu/codegen/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/BUILD @@ -48,12 +48,8 @@ cc_library( deps = [ ":builtin_fp16", ":builtin_pow", - "//xla/backends/cpu/runtime:msan_emulated_tls", - "//xla/service/cpu:cpu_runtime", - "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", "@llvm-project//llvm:Core", "@llvm-project//llvm:OrcJIT", "@llvm-project//llvm:OrcShared", @@ -128,21 +124,18 @@ cc_library( "//xla:util", "//xla:xla_proto_cc", "//xla/backends/cpu:target_machine_options", - "//xla/backends/cpu/runtime:msan_emulated_tls", "//xla/codegen:intrinsic_lib", "//xla/codegen/intrinsic", "//xla/codegen/intrinsic:intrinsic_compiler_lib", "//xla/service:hlo_module_config", "//xla/service/cpu:backend_config_proto_cc", "//xla/service/cpu:cpu_options", - "//xla/service/cpu:cpu_runtime", "//xla/service/cpu:executable_proto_cc", "//xla/service/llvm_ir:llvm_util", "//xla/tools:llvm_targets", # fixdeps: keep "//xla/tsl/platform:logging", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", - "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/log", diff --git a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc index 858223ff375ba1..97e056df2882c2 100644 --- a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc +++ b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc @@ -28,24 +28,19 @@ limitations under the License. #include #include -#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" -#include "absl/log/log.h" // IWYU pragma: keep #include "llvm/ADT/StringRef.h" #include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h" #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/CoreContainers.h" -#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" // IWYU pragma: keep (msan) #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h" #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h" #include "llvm/IR/DataLayout.h" #include "llvm/Support/Error.h" #include "xla/backends/cpu/codegen/builtin_fp16.h" #include "xla/backends/cpu/codegen/builtin_pow.h" -#include "xla/backends/cpu/runtime/msan_emulated_tls.h" -#include "xla/service/cpu/cpu_runtime.h" namespace xla::cpu { @@ -272,8 +267,9 @@ static Registry CreateRegistry() { #endif - registry[runtime::kMsanEmutlsGetAddressBridgeSymbolName] = - SymbolDef(__xla_cpu_runtime_emutls_get_address); +#ifdef MEMORY_SANITIZER + registry["__msan_unpoison"] = SymbolDef(__msan_unpoison); +#endif return registry; } @@ -284,55 +280,22 @@ static Registry CreateRegistry() { BuiltinDefinitionGenerator::BuiltinDefinitionGenerator( llvm::DataLayout data_layout) - : data_layout_(std::move(data_layout)) { -#ifdef ABSL_HAVE_MEMORY_SANITIZER - // Resolve MSan runtime functions (e.g. __msan_warning*) from the current - // process via dlsym. This is more future-proof than explicitly intercepting - // __msan_* functions; these functions do change between LLVM versions. - auto is_msan_symbol = [](const llvm::orc::SymbolStringPtr& name) { - return (*name).starts_with("__msan_"); - }; - auto generator = - llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( - data_layout_.getGlobalPrefix(), is_msan_symbol); - if (generator) { - process_generator_ = std::move(*generator); - } else { - LOG(WARNING) << "Failed to initialize dynamic library generator for MSan: " - << llvm::toString(generator.takeError()); - } -#endif -} + : data_layout_(std::move(data_layout)) {} llvm::Error BuiltinDefinitionGenerator::tryToGenerate( - llvm::orc::LookupState& ls, llvm::orc::LookupKind kind, - llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags flags, + llvm::orc::LookupState&, llvm::orc::LookupKind kind, + llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags, const llvm::orc::SymbolLookupSet& names) { llvm::orc::SymbolMap symbols; symbols.reserve(names.size()); -#ifdef ABSL_HAVE_MEMORY_SANITIZER - llvm::orc::SymbolLookupSet msan_names; -#endif - for (const auto& [name, name_flags] : names) { + for (const auto& [name, flags] : names) { if (auto symbol = ResolveBuiltinSymbol(data_layout_, *name)) { symbols[name] = *symbol; -#ifdef ABSL_HAVE_MEMORY_SANITIZER - } else if ((*name).starts_with("__msan_")) { - msan_names.add(name, name_flags); -#endif } } cantFail(jit_dylib.define(llvm::orc::absoluteSymbols(std::move(symbols)))); - -#ifdef ABSL_HAVE_MEMORY_SANITIZER - if (!msan_names.empty() && process_generator_) { - return process_generator_->tryToGenerate(ls, kind, jit_dylib, flags, - msan_names); - } -#endif - return llvm::Error::success(); } diff --git a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h index 23ea4174cc1553..7689e7b13e425c 100644 --- a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h +++ b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h @@ -16,12 +16,7 @@ limitations under the License. #ifndef XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_ #define XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_ -#include -#include - -#include "absl/base/config.h" // IWYU pragma: keep #include "llvm/ExecutionEngine/Orc/Core.h" -#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" // IWYU pragma: keep #include "llvm/IR/DataLayout.h" #include "llvm/Support/Error.h" @@ -49,9 +44,6 @@ class BuiltinDefinitionGenerator : public llvm::orc::DefinitionGenerator { private: llvm::DataLayout data_layout_; -#ifdef ABSL_HAVE_MEMORY_SANITIZER - std::unique_ptr process_generator_; -#endif }; } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc index 0a451a6807f33a..bda6062f77be13 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/backends/cpu/codegen/ir_compiler.h" #include -#include #include #include #include @@ -25,7 +24,6 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/call_once.h" -#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/nullability.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -40,25 +38,17 @@ limitations under the License. #include "llvm-c/Target.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Analysis/CGSCCPassManager.h" -#include "llvm/Analysis/GlobalsModRef.h" #include "llvm/Analysis/LoopAnalysisManager.h" #include "llvm/Analysis/RuntimeLibcallInfo.h" #include "llvm/Analysis/TargetLibraryInfo.h" #include "llvm/ExecutionEngine/ExecutionEngine.h" #include "llvm/ExecutionEngine/Orc/Mangling.h" -#include "llvm/IR/Attributes.h" -#include "llvm/IR/BasicBlock.h" -#include "llvm/IR/DataLayout.h" -#include "llvm/IR/DerivedTypes.h" -#include "llvm/IR/IRBuilder.h" -#include "llvm/IR/LLVMContext.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Metadata.h" #include "llvm/IR/PassManager.h" #include "llvm/IR/Verifier.h" #include "llvm/MC/MCContext.h" #include "llvm/Object/ObjectFile.h" -#include "llvm/Pass.h" #include "llvm/Passes/OptimizationLevel.h" #include "llvm/Passes/PassBuilder.h" #include "llvm/Passes/StandardInstrumentations.h" @@ -75,17 +65,14 @@ limitations under the License. #include "llvm/TargetParser/Triple.h" #include "llvm/Transforms/IPO/AlwaysInliner.h" #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h" -#include "llvm/Transforms/Instrumentation/MemorySanitizer.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" #include "xla/backends/cpu/codegen/polynomial_approximations.h" -#include "xla/backends/cpu/runtime/msan_emulated_tls.h" #include "xla/backends/cpu/target_machine_options.h" #include "xla/codegen/intrinsic/intrinsic.h" #include "xla/codegen/intrinsic/intrinsic_compiler_lib.h" #include "xla/codegen/intrinsic_lib.h" #include "xla/service/cpu/backend_config.pb.h" #include "xla/service/cpu/cpu_options.h" -#include "xla/service/cpu/cpu_runtime.h" #include "xla/service/hlo_module_config.h" #include "xla/service/llvm_ir/llvm_util.h" #include "xla/tsl/platform/logging.h" @@ -94,8 +81,6 @@ limitations under the License. namespace xla::cpu { -static constexpr char kNoSanitizeMemoryAttr[] = "no_sanitize_memory"; - namespace internal { static absl::once_flag targets_init; @@ -251,9 +236,9 @@ std::unique_ptr IrCompiler::Create( llvm::TargetOptions target_options, Options options, CompilationHooks hooks) { TargetMachineBuilder target_machine_builder = - IrCompiler::InferTargetMachineBuilder( - std::move(target_options), options.opt_level, - options.target_machine_options, options.msan_enabled); + IrCompiler::InferTargetMachineBuilder(std::move(target_options), + options.opt_level, + options.target_machine_options); return std::make_unique(target_machine_builder, std::move(options), std::move(hooks)); @@ -269,19 +254,14 @@ IrCompiler::IrCompiler(TargetMachineBuilder target_machine_builder, absl::StatusOr> IrCompiler::InferTargetMachine( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options, bool msan_enabled) { + const TargetMachineOptions& target_machine_options) { auto attrs_vec = target_machine_options.GetTargetMachineFeaturesVector(); llvm::SmallVector attrs(attrs_vec.begin(), attrs_vec.end()); - llvm::TargetOptions effective_target_options = target_options; - if (msan_enabled) { - effective_target_options.EmulatedTLS = true; - } - absl::call_once(internal::targets_init, &internal::InitializeTargets); std::unique_ptr target_machine( llvm::EngineBuilder() - .setTargetOptions(effective_target_options) + .setTargetOptions(target_options) .setOptLevel(opt_level) .selectTarget( /*TargetTriple=*/llvm::Triple(target_machine_options.triple()), @@ -299,10 +279,10 @@ IrCompiler::InferTargetMachine( IrCompiler::TargetMachineBuilder IrCompiler::InferTargetMachineBuilder( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options, bool msan_enabled) { - return [target_options, opt_level, target_machine_options, msan_enabled] { - return InferTargetMachine(target_options, opt_level, target_machine_options, - msan_enabled); + const TargetMachineOptions& target_machine_options) { + return [target_options, opt_level, target_machine_options] { + return InferTargetMachine(target_options, opt_level, + target_machine_options); }; } @@ -352,10 +332,6 @@ llvm::Expected> IrCompiler::operator()( } } - if (options_.msan_enabled) { - InjectMsanEmulatedTls(module); - } - std::unique_ptr mc_memory_buffer = EmitMachineCode(module, target_machine->get()); @@ -450,45 +426,12 @@ llvm::Error IrCompiler::RunIrPasses(llvm::Module& module, pb.registerLoopAnalyses(lam); pb.crossRegisterProxies(lam, fam, cgam, mam); - if (options_.msan_enabled) { - for (auto& function : module) { - if (!function.isDeclaration() && - !function.hasFnAttribute(kNoSanitizeMemoryAttr)) { - function.addFnAttr(llvm::Attribute::SanitizeMemory); - } - } - } - - pb.registerOptimizerLastEPCallback([&](llvm::ModulePassManager& mpm, - llvm::OptimizationLevel level, - llvm::ThinOrFullLTOPhase) { - if (options_.dfsan_enabled) { - mpm.addPass(llvm::DataFlowSanitizerPass(options_.dfsan_abi_list_files)); - } - - if (options_.msan_enabled) { - llvm::MemorySanitizerOptions msan_options( - options_.msan_track_origins, /*Recover=*/false, /*Kernel=*/false, - // Set eager checks to true. This is important to avoid msan flakes on - // KernelThunk's call frame pointer argument with AOT kernels: eager - // checks + nonnull/noundef annotations on the pointer make the AOT - // kernel never read the pointer's msan shadow memory. Without either, - // the AOT kernel would read shadow memory, but the host would not - // have written to it when compiled with eager checks enabled, which - // is the default in Clang. - // Note that if the host is built without eager checks, things still - // work: the host will write the call frame pointer's shadow memory, - // and the AOT kernel won't read it. For pointers flowing in the - // opposite direction (i.e. AOT -> host, such as the return pointer - // from a kernel) we do not mark them as noundef/nonnull, so we - // always write their corresponding shadow memory. - /*EagerChecks=*/true); - mpm.addPass(llvm::MemorySanitizerPass(msan_options)); - } - }); - llvm::ModulePassManager pm; + if (options_.dfsan_enabled) { + pm.addPass(llvm::DataFlowSanitizerPass(options_.dfsan_abi_list_files)); + } + llvm::OptimizationLevel opt_level = GetOptimizationLevel(options_); if (opt_level == llvm::OptimizationLevel::O0) { pm.addPass(pb.buildO0DefaultPipeline(opt_level)); @@ -583,54 +526,4 @@ IrCompiler::build_target_machine() const { return target_machine_builder_(); } -void IrCompiler::InjectMsanEmulatedTls(llvm::Module& module) const { - llvm::LLVMContext& ctx = module.getContext(); - const llvm::DataLayout& dl = module.getDataLayout(); - llvm::Type* void_ptr_ty = llvm::PointerType::get(ctx, 0); - - auto inject_selector = [&](llvm::StringRef name, MsanTlsSelector selector) { - new llvm::GlobalVariable( - module, void_ptr_ty, /*isConstant=*/true, - llvm::GlobalValue::InternalLinkage, - llvm::Constant::getIntegerValue( - void_ptr_ty, llvm::APInt(dl.getPointerSizeInBits(), - static_cast(selector))), - name); - }; - - inject_selector("__emutls_v.__msan_param_tls", MsanTlsSelector::kParamTls); - inject_selector("__emutls_v.__msan_retval_tls", MsanTlsSelector::kRetvalTls); - inject_selector("__emutls_v.__msan_va_arg_tls", MsanTlsSelector::kVaArgTls); - inject_selector("__emutls_v.__msan_va_arg_overflow_size_tls", - MsanTlsSelector::kVaArgOverflowSizeTls); - inject_selector("__emutls_v.__msan_param_origin_tls", - MsanTlsSelector::kParamOriginTls); - inject_selector("__emutls_v.__msan_retval_origin_tls", - MsanTlsSelector::kRetvalOriginTls); - inject_selector("__emutls_v.__msan_va_arg_origin_tls", - MsanTlsSelector::kVaArgOriginTls); - inject_selector("__emutls_v.__msan_origin_tls", MsanTlsSelector::kOriginTls); - - llvm::FunctionType* emutls_get_addr_type = - llvm::FunctionType::get(void_ptr_ty, void_ptr_ty, /*isVarArg=*/false); - llvm::Function* emutls_get_addr_fn = llvm::cast( - module.getOrInsertFunction("__emutls_get_address", emutls_get_addr_type) - .getCallee()); - emutls_get_addr_fn->setLinkage(llvm::GlobalValue::InternalLinkage); - emutls_get_addr_fn->addFnAttr(kNoSanitizeMemoryAttr); - - llvm::FunctionCallee bridge_fn = module.getOrInsertFunction( - runtime::kMsanEmutlsGetAddressBridgeSymbolName, emutls_get_addr_type); - - llvm::BasicBlock* entry = - llvm::BasicBlock::Create(ctx, "entry", emutls_get_addr_fn); - llvm::IRBuilder<> builder(entry); - // LLVM's emutls calls __emutls_get_address with the *address* of the control - // variable. We need to dereference it to get the selector value. - auto* control_val = - builder.CreateLoad(void_ptr_ty, emutls_get_addr_fn->getArg(0)); - auto* call = builder.CreateCall(bridge_fn, {control_val}); - builder.CreateRet(call); -} - } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h index ce67dc595800fe..c170bc0755f9c7 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h @@ -23,10 +23,6 @@ limitations under the License. #include #include -#include "absl/base/config.h" // IWYU pragma: keep -#ifdef ABSL_HAVE_MEMORY_SANITIZER -#include -#endif #include "absl/base/thread_annotations.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" @@ -79,19 +75,6 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { bool disable_loop_unrolling = false; bool disable_platform_dependent_math = false; - // This should be the _only_ place where an #ifdef determines whether the - // generated code is msan-instrumented. This ensures an uninstrumented - // compiler can produce msan-instrumented AOT objects, and viceversa. -#ifdef ABSL_HAVE_MEMORY_SANITIZER - bool msan_enabled = true; - // Level of MSan origin tracking (0 = off, 1 = basic, 2 = full origins). - // In JIT compilation, this defaults to the host's origin tracking level. - int msan_track_origins = __msan_get_track_origins(); -#else - bool msan_enabled = false; - int msan_track_origins = 0; -#endif - bool dfsan_enabled = false; std::vector dfsan_abi_list_files; }; @@ -115,16 +98,14 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { static absl::StatusOr> InferTargetMachine(const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options, - bool msan_enabled = false); + const TargetMachineOptions& target_machine_options); // Returns a target machine builder that uses `InferTargetMachine` defined // above to infer the target machine for the given options. static TargetMachineBuilder InferTargetMachineBuilder( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options, - bool msan_enabled = false); + const TargetMachineOptions& target_machine_options); // Compiles a `module` to an ObjectFile. llvm::Expected> operator()( @@ -159,10 +140,6 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { // races when calling user provided compilation hooks. absl::Mutex mutex_; CompilationHooks hooks_ ABSL_GUARDED_BY(mutex_); - - // Injects MSAN emulated TLS symbols into the module. This is needed for - // supporting MSAN in JIT'ed and AOT'ed code. - void InjectMsanEmulatedTls(llvm::Module& module) const; }; } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc index 1c87c5c122ab8f..29a789f71d1dca 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc @@ -294,13 +294,12 @@ TEST(IrCompilerTest, EmitIntrinsicCall) { auto context = std::make_unique(); IrCompiler::CompilationHooks compilation_hooks; - IrCompiler::Options options{/*opt_level=*/llvm::CodeGenOptLevel::Aggressive, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}; - options.msan_enabled = false; // Avoid msan interception of memcpy. - - std::unique_ptr ir_compiler = - IrCompiler::Create(llvm::TargetOptions(), options, compilation_hooks); + std::unique_ptr ir_compiler = IrCompiler::Create( + llvm::TargetOptions(), + IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::Aggressive, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}, + compilation_hooks); TF_ASSERT_OK_AND_ASSIGN(auto ir_module, ParseModule(*context, kMemcpyCall, kModuleName)); @@ -354,40 +353,6 @@ INSTANTIATE_TEST_SUITE_P(IrCompilerParameterizedTestInstantiation, ::testing::Values("x86_64-grtev4-linux-gnu", "aarch64-unknown-linux-gnu")); -TEST(IrCompilerTest, MemorySanitizerTrackOrigins) { - auto context = std::make_unique(); - IrCompiler::CompilationHooks compilation_hooks; - - TargetMachineOptions target_machine_options(kTargetTripleForHost, - kTargetCpuForHost, ""); - - IrCompiler::Options options{ - /*opt_level=*/llvm::CodeGenOptLevel::None, - /*optimize_for_size=*/false, - target_machine_options, - }; - options.msan_enabled = true; - options.msan_track_origins = 2; - - std::unique_ptr ir_compiler = - IrCompiler::Create(llvm::TargetOptions(), options, compilation_hooks); - - ASSERT_OK_AND_ASSIGN(auto ir_module, - ParseModule(*context, kUnoptimizedIr, "test_module")); - - ASSERT_OK_AND_ASSIGN(auto target_machine, - ir_compiler->build_target_machine()); - - ir_module->setDataLayout(target_machine->createDataLayout()); - ir_module->setTargetTriple(target_machine->getTargetTriple()); - cantFail((*ir_compiler)(*ir_module)); - - auto ir = llvm_ir::DumpToString(ir_module.get()); - EXPECT_THAT(ir, HasSubstr("__msan_track_origins = weak_odr constant i32 2")); - EXPECT_THAT(ir, HasSubstr("__emutls_v.__msan_param_tls")); - EXPECT_THAT(ir, HasSubstr("@__emutls_get_address")); -} - } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc index 2c37fefcac7d26..ad3bf50ddf1acc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc @@ -87,18 +87,6 @@ static absl::StatusOr ParseModule( return llvm::orc::ThreadSafeModule(std::move(m), context); } -// Creates an IrCompiler for testing. We explicitly disable MSan instrumentation -// because unit tests in this file compile raw LLVM IR snippets without linking -// the XLA CPU runtime or BuiltinDefinitionGenerator. -static std::unique_ptr CreateTestIrCompiler() { - IrCompiler::Options options{/*opt_level=*/llvm::CodeGenOptLevel::None, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}; - options.msan_enabled = false; - return IrCompiler::Create(llvm::TargetOptions(), std::move(options), - IrCompiler::CompilationHooks()); -} - TEST(JitCompilerTest, Compile) { auto context = std::make_unique(); llvm::orc::ThreadSafeContext tsc(std::move(context)); @@ -114,7 +102,12 @@ TEST(JitCompilerTest, Compile) { thread_pool.Schedule(std::move(task)); }; - std::unique_ptr ir_compiler = CreateTestIrCompiler(); + std::unique_ptr ir_compiler = IrCompiler::Create( + llvm::TargetOptions(), + IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}, + IrCompiler::CompilationHooks()); TF_ASSERT_OK_AND_ASSIGN( auto compiler, @@ -208,7 +201,12 @@ TEST(JitCompilerTest, ExternalDefinitionGenerator) { return std::make_unique(); }; - std::unique_ptr ir_compiler = CreateTestIrCompiler(); + std::unique_ptr ir_compiler = IrCompiler::Create( + llvm::TargetOptions(), + IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}, + IrCompiler::CompilationHooks()); TF_ASSERT_OK_AND_ASSIGN( auto compiler, @@ -302,7 +300,12 @@ TEST(JitCompilerTest, CompileWithHighAlignment) { llvm::orc::ThreadSafeContext tsc(std::move(context)); JitCompiler::Options options; - std::unique_ptr ir_compiler = CreateTestIrCompiler(); + std::unique_ptr ir_compiler = IrCompiler::Create( + llvm::TargetOptions(), + IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}, + IrCompiler::CompilationHooks()); TF_ASSERT_OK_AND_ASSIGN( auto compiler, diff --git a/third_party/xla/xla/backends/cpu/nanort/BUILD b/third_party/xla/xla/backends/cpu/nanort/BUILD index fed38e3c85491d..f67ac41f0addd5 100644 --- a/third_party/xla/xla/backends/cpu/nanort/BUILD +++ b/third_party/xla/xla/backends/cpu/nanort/BUILD @@ -94,7 +94,6 @@ xla_cc_test( "//xla/tsl/platform:test_benchmark", "//xla/tsl/platform:test_main", "@com_google_absl//absl/base", - "@com_google_absl//absl/base:config", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc index 833d89c18bf12c..ec23293b9a8a8b 100644 --- a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc +++ b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc @@ -26,7 +26,6 @@ limitations under the License. #include #include "absl/base/casts.h" -#include "absl/base/config.h" // IWYU pragma: keep #include "absl/container/inlined_vector.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" @@ -63,10 +62,6 @@ limitations under the License. #include "xla/xla_data.pb.h" #include "tsl/platform/casts.h" -#ifdef ABSL_HAVE_MEMORY_SANITIZER -#include -#endif - #define EIGEN_USE_THREADS #include "Eigen/ThreadPool" @@ -487,54 +482,6 @@ TEST_P(NanoRtClientTest, ProgramShapeKeepsLayout) { absl::Span({0, 1})); } -TEST_P(NanoRtClientTest, MsanTracksPoisonThroughKernel) { -#ifndef ABSL_HAVE_MEMORY_SANITIZER - GTEST_SKIP() << "This test requires an MSan build"; -#else - const char* kModuleStr = R"( - HloModule msan_shadow_test - - ENTRY e { - p0 = f32[4] parameter(0) - p1 = f32[4] parameter(1) - ROOT sum = f32[4] add(p0, p1) - } - )"; - - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnUnverifiedModule(kModuleStr)); - XlaComputation computation(module->ToProto()); - - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr executable, - GetExecutable(computation, GetParam())); - - alignas(cpu::MinAlign()) float p0[4]; - p0[0] = 1.0f; - p0[1] = 2.0f; - // p0[2], p0[3] intentionally uninitialized (poisoned under MSan). - - alignas(cpu::MinAlign()) float p1[4] = {10.0f, 20.0f, 30.0f, 40.0f}; - alignas(cpu::MinAlign()) float result[4] = {}; - - Arguments arguments = {{p0, 4}, {p1, 4}}; - Results results = {{result, 4}}; - - auto event = executable->Execute(arguments, results, {}); - tsl::BlockUntilReady(event); - ASSERT_TRUE(event.IsConcrete()); - - EXPECT_EQ(__msan_test_shadow(&result[0], sizeof(float)), -1) - << "result[0] should be initialized (unpoisoned)"; - EXPECT_EQ(__msan_test_shadow(&result[1], sizeof(float)), -1) - << "result[1] should be initialized (unpoisoned)"; - - EXPECT_GE(__msan_test_shadow(&result[2], sizeof(float)), 0) - << "result[2] should be poisoned (p0[2] was uninitialized)"; - EXPECT_GE(__msan_test_shadow(&result[3], sizeof(float)), 0) - << "result[3] should be poisoned (p0[3] was uninitialized)"; -#endif -} - INSTANTIATE_TEST_SUITE_P(NanoRtClientTestSuite, NanoRtClientTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { diff --git a/third_party/xla/xla/backends/cpu/runtime/BUILD b/third_party/xla/xla/backends/cpu/runtime/BUILD index 130bc7b2223b0e..fc20ecd3f5fe5a 100644 --- a/third_party/xla/xla/backends/cpu/runtime/BUILD +++ b/third_party/xla/xla/backends/cpu/runtime/BUILD @@ -47,7 +47,6 @@ filegroup( "dot_lib_f64.cc", "dot_lib_s32.cc", "dot_lib_s8.cc", - "msan_emulated_tls.cc", "rng_state_lib.cc", "sort_lib.cc", ], @@ -60,7 +59,6 @@ filegroup( "convolution_lib.h", "dot_lib.h", "kernel_c_api.h", - "msan_emulated_tls.h", "rng_state_lib.h", "sort_lib.h", "work_queue.h", @@ -168,20 +166,6 @@ cc_library( ], ) -cc_library( - name = "msan_emulated_tls", - srcs = ["msan_emulated_tls.cc"], - hdrs = ["msan_emulated_tls.h"], - deps = [ - "@com_google_absl//absl/base:config", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/log:check", - ], - # Must always link because we might compile/run programs with msan enabled, - # regardless of whether the host compiler is built with msan or not. - alwayslink = True, -) - tf_proto_library( name = "thunk_proto", srcs = ["thunk.proto"], diff --git a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc deleted file mode 100644 index c2079d29c03eff..00000000000000 --- a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc +++ /dev/null @@ -1,80 +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. -==============================================================================*/ - -#include "xla/backends/cpu/runtime/msan_emulated_tls.h" - -#include // IWYU pragma: keep - -#include "absl/base/attributes.h" // IWYU pragma: keep -#include "absl/base/config.h" // IWYU pragma: keep -#include "absl/base/optimization.h" // IWYU pragma: keep -#include "absl/log/check.h" // IWYU pragma: keep - -#ifdef ABSL_HAVE_MEMORY_SANITIZER -extern "C" { -// Mark these initial-exec as compiler-rt does. -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_param_tls[]; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_param_origin_tls[]; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_retval_tls[]; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_retval_origin_tls; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_va_arg_tls[]; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_va_arg_origin_tls[]; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uintptr_t - __msan_va_arg_overflow_size_tls; -extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_origin_tls; -} -#endif // ABSL_HAVE_MEMORY_SANITIZER - -static_assert( - static_cast(xla::cpu::MsanTlsSelector::kParamTls) == 1 && - static_cast(xla::cpu::MsanTlsSelector::kOriginTls) == 8, - "MsanTlsSelector must remain a contiguous 1..8 range"); - -extern "C" { - -ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void* __xla_cpu_runtime_emutls_get_address( - void* control) { -#ifdef ABSL_HAVE_MEMORY_SANITIZER - using xla::cpu::MsanTlsSelector; - // The control argument is already the selector value (not a pointer to it) - // because the internal __emutls_get_address wrapper in the LLVM module - // dereferences the selector before calling this bridge. - uintptr_t selector = reinterpret_cast(control); - DCHECK_GE(selector, static_cast(MsanTlsSelector::kParamTls)); - DCHECK_LE(selector, static_cast(MsanTlsSelector::kOriginTls)); - switch (static_cast(selector)) { - case MsanTlsSelector::kParamTls: - return __msan_param_tls; - case MsanTlsSelector::kRetvalTls: - return __msan_retval_tls; - case MsanTlsSelector::kVaArgTls: - return __msan_va_arg_tls; - case MsanTlsSelector::kVaArgOverflowSizeTls: - return &__msan_va_arg_overflow_size_tls; - case MsanTlsSelector::kParamOriginTls: - return __msan_param_origin_tls; - case MsanTlsSelector::kRetvalOriginTls: - return &__msan_retval_origin_tls; - case MsanTlsSelector::kVaArgOriginTls: - return __msan_va_arg_origin_tls; - case MsanTlsSelector::kOriginTls: - return &__msan_origin_tls; - } - ABSL_UNREACHABLE(); -#endif // ABSL_HAVE_MEMORY_SANITIZER - return nullptr; -} - -} // extern "C" diff --git a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h deleted file mode 100644 index 5a592629ebaa9a..00000000000000 --- a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h +++ /dev/null @@ -1,47 +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_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ -#define XLA_BACKENDS_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ - -#include - -#include "absl/base/attributes.h" - -extern "C" { -// Returns the address of the host's MSAN TLS variables. -// See https://github.com/google/sanitizers/wiki/MemorySanitizerJIT. -ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void* __xla_cpu_runtime_emutls_get_address( - void* control); -} // extern "C" - -namespace xla::cpu { - -// Selectors for __emutls_get_address. -// All of these are needed for msan and msan-track-origins. -enum class MsanTlsSelector : uintptr_t { - kParamTls = 1, - kRetvalTls = 2, - kVaArgTls = 3, - kVaArgOverflowSizeTls = 4, - kParamOriginTls = 5, - kRetvalOriginTls = 6, - kVaArgOriginTls = 7, - kOriginTls = 8, -}; - -} // namespace xla::cpu - -#endif // XLA_BACKENDS_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 714dadf77ac100..56f9a5697b38ce 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1961,9 +1961,8 @@ cc_library( hdrs = ["collective_kernel_thunk.h"], deps = [ ":all_reduce", - ":collective_cliques", + ":collective_kernel_api", ":collective_kernel_thunk_proto_cc", - ":collective_memory", ":collective_params", ":collective_thunk", ":collective_thunk_proto_cc", @@ -1975,9 +1974,7 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/backends/gpu/collectives:gpu_clique_key", - "//xla/backends/gpu/collectives:gpu_communicator", "//xla/core/collectives:rank_id", - "//xla/core/collectives:symmetric_memory", "//xla/runtime:buffer_use", "//xla/runtime:device_id", "//xla/service:buffer_assignment", @@ -1995,7 +1992,6 @@ cc_library( "//xla/stream_executor/gpu:all_reduce_kernel", "//xla/stream_executor/gpu:collective_kernel_metadata", "//xla/tsl/util:safe_reinterpret_cast", - "//xla/tsl/util:tied_ref", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", @@ -4668,6 +4664,7 @@ cc_library( "//xla/tsl/concurrency:async_value", "//xla/tsl/framework:allocator", "//xla/tsl/platform:env", + "//xla/tsl/util:unique_any", "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", @@ -4716,18 +4713,21 @@ xla_test( "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream_executor_address_allocator", "//xla/stream_executor:stream_executor_h", + "//xla/tests:hlo_test_base", "//xla/tests:literal_test_util", + "//xla/tests:xla_internal_test_main", "//xla/tsl/concurrency:async_value", - "//xla/tsl/lib/core:status_test_util", + "//xla/tsl/platform:test_benchmark", "//xla/tsl/util/proto:proto_matchers", "@com_google_absl//absl/base", "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@com_google_googletest//:gtest_main", + "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc index 89204ae5635988..dce95619458e2c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc @@ -148,5 +148,41 @@ size_t GetMultiGpuBarrierSignalBufferSize() { size_t GetMultiGpuBarrierSignalValueSize() { return sizeof(uint32_t); } +absl::StatusOr> CollectParamToPeers( + const GpuCliqueKey& clique_key, RankId rank, + stream_executor::Stream* stream, + std::vector parameters) { + std::vector param_to_peers_ptrs; + + size_t num_parameters = parameters.size(); + // Exchange device parameters with all ranks in the clique. + ABSL_ASSIGN_OR_RETURN( + auto device_parameters, + GpuCliqueRendezvous::Join(clique_key, rank, std::move(parameters))); + + // Collect pointers to device buffers from all participating ranks. + param_to_peers_ptrs.reserve(num_parameters * clique_key.num_devices()); + + absl::flat_hash_map> + peer_to_parameters(clique_key.num_devices()); + + using DeviceParameters = std::vector; + + for (auto peer = RankId(0); peer < RankId(clique_key.num_devices()); ++peer) { + ABSL_ASSIGN_OR_RETURN(const DeviceParameters& peer_parameters, + device_parameters->at(peer)); + peer_to_parameters[peer.value()] = std::move(peer_parameters); + } + + for (int parameter = 0; parameter < num_parameters; ++parameter) { + for (int peer = 0; peer < clique_key.num_devices(); ++peer) { + param_to_peers_ptrs.push_back( + peer_to_parameters[peer][parameter].opaque()); + } + } + + return param_to_peers_ptrs; +} + } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h index cd601d10610dd9..ff4bc5101d7593 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h @@ -21,6 +21,8 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "xla/backends/gpu/collectives/gpu_clique_key.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/symmetric_memory.h" #include "xla/stream_executor/device_address.h" @@ -55,6 +57,13 @@ size_t GetMultiGpuBarrierSignalBufferSize(); // Returns the size of the barrier signal value in bytes. size_t GetMultiGpuBarrierSignalValueSize(); +// Collect the pointers to the parameters at the peer devices. +// The size of the returned vector is num_parameters * num_devices. +absl::StatusOr> CollectParamToPeers( + const GpuCliqueKey& clique_key, RankId rank, + stream_executor::Stream* stream, + std::vector parameters); + } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc index e48d12ca03c5ef..5294bcb0bc783c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc @@ -35,18 +35,15 @@ limitations under the License.*/ #include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/gpu_clique_key.h" -#include "xla/backends/gpu/collectives/gpu_communicator.h" #include "xla/backends/gpu/runtime/all_reduce.h" -#include "xla/backends/gpu/runtime/collective_cliques.h" +#include "xla/backends/gpu/runtime/collective_kernel_api.h" #include "xla/backends/gpu/runtime/collective_kernel_thunk.pb.h" -#include "xla/backends/gpu/runtime/collective_memory.h" #include "xla/backends/gpu/runtime/collective_params.h" #include "xla/backends/gpu/runtime/collective_thunk.h" #include "xla/backends/gpu/runtime/collective_thunk.pb.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/core/collectives/rank_id.h" -#include "xla/core/collectives/symmetric_memory.h" #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" @@ -65,7 +62,6 @@ limitations under the License.*/ #include "xla/stream_executor/stream.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/util/safe_reinterpret_cast.h" -#include "xla/tsl/util/tied_ref.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -459,37 +455,8 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { const size_t num_parameters = parameters.size(); const size_t param_to_peers_ptrs_size_bytes = num_parameters * clique_key.num_devices() * sizeof(uint64_t); - TF_RET_CHECK(params.collective_params != nullptr) - << "Collective params must not be null in " - "CollectiveKernelThunk::Initialize"; - TF_RET_CHECK(params.collective_cliques != nullptr) - << "Collective cliques must not be null in " - "CollectiveKernelThunk::Initialize"; - TF_RET_CHECK(params.collective_memory != nullptr) - << "Collective memory must not be null in " - "CollectiveKernelThunk::Initialize"; - - ABSL_ASSIGN_OR_RETURN(GpuCommunicator * comm, - params.collective_cliques->GetComm(clique_key, *rank)); - - if (memory_state->scratch_symmetric_memories.empty()) { - memory_state->scratch_symmetric_memories.reserve( - memory_state->scratch_allocations.size()); - for (size_t i = 0; i < memory_state->scratch_allocations.size(); ++i) { - se::DeviceAddressBase addr = - memory_state->scratch_allocations[i].address(); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr symmetric_memory, - comm->CreateSymmetricMemory(addr)); - ABSL_ASSIGN_OR_RETURN(tsl::TiedRef tied_symmetric_memory, - params.collective_cliques->Tie( - clique_key, std::move(symmetric_memory))); - memory_state->scratch_symmetric_memories.push_back( - std::move(tied_symmetric_memory)); - } - } - std::vector multimem_addresses; - if (RequiresMultimem(kernel_spec_)) { + if (RequiresMultimem(kernel_spec_) && params.collective_memory != nullptr) { multimem_addresses.resize(num_parameters, nullptr); for (size_t i = 0; i < num_parameters; ++i) { auto [mmem, offset] = params.collective_memory->FindSymmetricMemory( @@ -502,30 +469,9 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { } } } - - static constexpr auto is_multimem_buffer = - [](const IoBufferSpec& spec) -> bool { return spec.requires_multimem; }; - int32_t scratch_buffers_index = - absl::c_count_if(kernel_spec_.input_buffer_specs, is_multimem_buffer) + - absl::c_count_if(kernel_spec_.output_buffer_specs, is_multimem_buffer); - std::vector param_to_peers_ptrs(num_parameters * - clique_key.num_devices()); - for (size_t i = scratch_buffers_index; i < num_parameters; ++i) { - const size_t scratch_index = i - scratch_buffers_index; - auto sym_mem = - memory_state->scratch_symmetric_memories[scratch_index].Lock(); - TF_RET_CHECK(sym_mem != nullptr) - << "Symmetric memory for scratch buffer " << scratch_index - << " is no longer valid"; - for (int device_rank = 0; device_rank < clique_key.num_devices(); - ++device_rank) { - ABSL_ASSIGN_OR_RETURN(se::DeviceAddressBase peer_address, - sym_mem->peer_addr(RankId(device_rank))); - const size_t parameter_offset = i * clique_key.num_devices(); - param_to_peers_ptrs[parameter_offset + device_rank] = - peer_address.opaque(); - } - } + ABSL_ASSIGN_OR_RETURN(std::vector param_to_peers_ptrs, + CollectParamToPeers(clique_key, state->rank, params.stream, + std::move(parameters))); const size_t multimem_size_bytes = multimem_addresses.size() * sizeof(void*); state->metadata = params.executor->Allocate( diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h index 760b3643accba4..0316d660cf4b81 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h @@ -37,7 +37,6 @@ limitations under the License.*/ #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/backends/gpu/runtime/traced_command.h" #include "xla/core/collectives/rank_id.h" -#include "xla/core/collectives/symmetric_memory.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/stream_executor/device_address.h" @@ -45,7 +44,6 @@ limitations under the License.*/ #include "xla/stream_executor/gpu/all_reduce_kernel.h" #include "xla/stream_executor/kernel.h" #include "xla/stream_executor/stream.h" -#include "xla/tsl/util/tied_ref.h" namespace xla::gpu { @@ -134,7 +132,6 @@ class CollectiveKernelThunk : public TracedCommand { // Per-executor scratch memory. struct StreamMemory { std::vector scratch_allocations; - std::vector> scratch_symmetric_memories; }; // Per-executor state that needs to be synchronized for access. diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index ef1f736b916900..ec65e9093d7f75 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -372,7 +372,6 @@ absl::StatusOr RunCollectiveKernelThunk( initialize_params.stream = stream.get(); initialize_params.buffer_allocations = &buffer_allocations; initialize_params.collective_params = &collective_params; - initialize_params.collective_cliques = &collective_cliques; initialize_params.src = {kKernelSource}; initialize_params.collective_memory = &collective_memory; @@ -478,7 +477,7 @@ TEST(CollectiveKernelThunkTest, MultiprocessTest) { /*is_multimem_enabled=*/false, /*use_ptx=*/true); EXPECT_THAT(RunCollectiveKernelThunkOnDevices(metadata, /*emulate_multiprocess=*/true), - StatusIs(absl::StatusCode::kNotFound)); + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(CollectiveKernelThunkTest, BufferUses) { @@ -616,29 +615,19 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) { &allocations1}; ASSERT_OK(collective_kernel_thunk->Prepare(prepare_params)); - CollectiveMemoryCache collective_memory_cache; - ASSERT_OK_AND_ASSIGN( - CollectiveCliques collective_cliques, - AcquireCollectiveCliques(collective_params, clique_requests)); - ASSERT_OK_AND_ASSIGN( - CollectiveMemory collective_memory, - AcquireCollectiveMemory(collective_params, collective_cliques, - memory_requests, collective_memory_cache)); - Thunk::InitializeParams initialize_params; initialize_params.executor = executor; initialize_params.stream = stream.get(); initialize_params.buffer_allocations = &allocations1; initialize_params.collective_params = &collective_params; - initialize_params.collective_cliques = &collective_cliques; initialize_params.src.text = kKernelSource; - initialize_params.collective_memory = &collective_memory; ASSERT_OK(collective_kernel_thunk->Initialize(initialize_params)); ASSERT_OK(stream->BlockHostUntilDone()); Thunk::ExecuteParams params1 = Thunk::ExecuteParams::Create( run_options, allocations1, stream.get(), trace_stream.get(), - &collective_params, &collective_cliques, &collective_memory); + &collective_params, /*collective_cliques=*/nullptr, + /*collective_memory=*/nullptr); CommandStateManager state; Command::RecordParams record_params = {state}; @@ -659,7 +648,8 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) { BufferAllocations updated_allocations({src2, dst2}, 0, nullptr); Thunk::ExecuteParams params2 = Thunk::ExecuteParams::Create( run_options, updated_allocations, stream.get(), trace_stream.get(), - &collective_params, &collective_cliques, &collective_memory); + &collective_params, /*collective_cliques=*/nullptr, + /*collective_memory=*/nullptr); std::vector updated_allocs = {0, 1}; Command::RecordParams update_record_params = {state, std::move(updated_allocs)}; diff --git a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc index 140ac68c07f291..3c10a01d02aa35 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc @@ -61,6 +61,7 @@ limitations under the License. #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/threadpool.h" +#include "xla/tsl/util/unique_any.h" #include "xla/util.h" #include "tsl/platform/cpu_info.h" #include "tsl/profiler/lib/traceme.h" @@ -116,20 +117,12 @@ bool CompareShapesIgnoringMemorySpace(const Shape& shape1, class HostExecuteCallFrame { public: - static absl::StatusOr Create( - se::Stream* device_to_host_stream, se::Stream* host_to_device_stream, + static absl::StatusOr> Create( + se::Stream* host_to_device_stream, const BufferAllocations* buffer_allocations, HostOffloadingAllocator& allocator, absl::Span args, absl::Span results, const ProgramShape& program_shape); - absl::Span> parameters() const { - return parameters_; - } - const ShapeTree& result() const { return result_; } - - absl::Status PublishResult() &&; - - protected: HostExecuteCallFrame( se::Stream* host_to_device_stream, const BufferAllocations* buffer_allocations, @@ -138,13 +131,25 @@ class HostExecuteCallFrame { absl::Span result_slices, std::vector> buffers); + absl::Span> parameters() const { + return parameters_; + } + const ShapeTree& result() const { return result_; } + + absl::Status CopyArguments(se::Stream* device_to_host_stream, + const BufferAllocations* buffer_allocations, + absl::Span args); + + static absl::Status PublishResult( + std::shared_ptr self, + const BufferAllocations* buffer_allocations); + + private: static absl::Status ValidateArgsAndResults( absl::Span args, absl::Span results, const ProgramShape& program_shape); - private: se::Stream* host_to_device_stream_; - const BufferAllocations* buffer_allocations_; std::vector> parameters_; ShapeTree result_; @@ -212,11 +217,13 @@ absl::Status HostExecuteCallFrame::ValidateArgsAndResults( return absl::OkStatus(); } -absl::StatusOr HostExecuteCallFrame::Create( - se::Stream* device_to_host_stream, se::Stream* host_to_device_stream, - const BufferAllocations* buffer_allocations, - HostOffloadingAllocator& allocator, absl::Span args, - absl::Span results, const ProgramShape& program_shape) { +absl::StatusOr> +HostExecuteCallFrame::Create(se::Stream* host_to_device_stream, + const BufferAllocations* buffer_allocations, + HostOffloadingAllocator& allocator, + absl::Span args, + absl::Span results, + const ProgramShape& program_shape) { tsl::profiler::TraceMe trace("HostExecuteCallFrame::Create"); ABSL_RETURN_IF_ERROR(ValidateArgsAndResults(args, results, program_shape)); @@ -230,8 +237,7 @@ absl::StatusOr HostExecuteCallFrame::Create( "HostExecuteCallFrame::Create Allocating Args"); for (const auto& [slice, shape] : args) { auto buffer_allocation = buffer_allocations->GetDeviceAddress(slice); - if (IsBufferOnDevice(device_to_host_stream, buffer_allocation.opaque())) { - // Copy device memory to host memory. + if (IsBufferOnDevice(host_to_device_stream, buffer_allocation.opaque())) { ABSL_ASSIGN_OR_RETURN( buffers.emplace_back(), allocator.AllocateTransferBuffer(ShapeUtil::ByteSizeOf(shape))); @@ -239,10 +245,6 @@ absl::StatusOr HostExecuteCallFrame::Create( parameters.push_back(ShapeTree( shape, HostOffloadingBuffer(buffers.back()->untyped_data(), buffers.back()->size_bytes()))); - - ABSL_RETURN_IF_ERROR(device_to_host_stream->Memcpy( - buffers.back()->untyped_data(), buffer_allocation, - buffers.back()->size_bytes())); } else { // We don't allocate as buffer is already in host memory. parameters.push_back(ShapeTree( @@ -277,9 +279,27 @@ absl::StatusOr HostExecuteCallFrame::Create( } } - return HostExecuteCallFrame(host_to_device_stream, buffer_allocations, - std::move(parameters), std::move(result), - std::move(results), std::move(buffers)); + return std::make_shared( + host_to_device_stream, buffer_allocations, std::move(parameters), + std::move(result), std::move(results), std::move(buffers)); +} + +absl::Status HostExecuteCallFrame::CopyArguments( + se::Stream* device_to_host_stream, + const BufferAllocations* buffer_allocations, absl::Span args) { + tsl::profiler::TraceMe trace("HostExecuteCallFrame::CopyArguments"); + int i = 0; + for (const auto& [slice, shape] : args) { + auto buffer_allocation = buffer_allocations->GetDeviceAddress(slice); + if (!IsBufferOnDevice(device_to_host_stream, buffer_allocation.opaque())) { + continue; + } + ABSL_RETURN_IF_ERROR(device_to_host_stream->Memcpy( + allocated_buffers_[i]->untyped_data(), buffer_allocation, + allocated_buffers_[i]->size_bytes())); + i++; + } + return absl::OkStatus(); } HostExecuteCallFrame::HostExecuteCallFrame( @@ -290,37 +310,36 @@ HostExecuteCallFrame::HostExecuteCallFrame( absl::Span result_slices, std::vector> buffers) : host_to_device_stream_(host_to_device_stream), - buffer_allocations_(buffer_allocations), parameters_(std::move(parameters)), result_(std::move(result)), result_slices_(result_slices), allocated_buffers_(std::move(buffers)) {} -absl::Status HostExecuteCallFrame::PublishResult() && { +absl::Status HostExecuteCallFrame::PublishResult( + std::shared_ptr self, + const BufferAllocations* buffer_allocations) { tsl::profiler::TraceMe trace("HostExecuteCallFrame::PublishResult"); size_t result_leaf_index = 0; - for (const auto& [index, buffer] : result_.leaves()) { - auto result_buffer = buffer_allocations_->GetDeviceAddress( - result_slices_[result_leaf_index++].slice); - if (!IsBufferOnDevice(host_to_device_stream_, result_buffer.opaque())) { + for (const auto& [index, buffer] : self->result_.leaves()) { + auto result_buffer = buffer_allocations->GetDeviceAddress( + self->result_slices_[result_leaf_index++].slice); + if (!IsBufferOnDevice(self->host_to_device_stream_, + result_buffer.opaque())) { // No need to copy result since the result is expected to be in host // memory and should match the buffer used for execution. CHECK(result_buffer.opaque() == buffer.opaque_base()); continue; } - auto shape = ShapeUtil::GetSubshape(result_.shape(), index); - ABSL_RETURN_IF_ERROR(host_to_device_stream_->Memcpy( + auto shape = ShapeUtil::GetSubshape(self->result_.shape(), index); + ABSL_RETURN_IF_ERROR(self->host_to_device_stream_->Memcpy( &result_buffer, buffer.opaque_base(), buffer.size_in_bytes())); } - // Move the backing buffers (allocated_buffers_) to the callback to ensure - // that they are only destroyed after the memory copies are done. - ABSL_RETURN_IF_ERROR(host_to_device_stream_->DoHostCallbackWithStatus( - [buffers = std::move(allocated_buffers_)]() { - return absl::OkStatus(); - })); - + // Store ref to the backing buffers (allocated_buffers_) to the callback to + // ensure that they are only destroyed after the memory copies are done. + ABSL_RETURN_IF_ERROR(self->host_to_device_stream_->DoHostCallbackWithStatus( + [self = std::move(self)]() { return absl::OkStatus(); })); return absl::OkStatus(); } @@ -371,8 +390,6 @@ HostExecuteAsyncEvents::ExtractEvent(se::StreamExecutor* executor, return event; } -// HostExecuteStartThunk - absl::StatusOr> HostExecuteStartThunk::Create( Thunk::ThunkInfo thunk_info, @@ -527,6 +544,15 @@ absl::Status HostExecuteStartThunk::Initialize(const InitializeParams& params) { } }); + ABSL_ASSIGN_OR_RETURN(std::shared_ptr call_frame, + HostExecuteCallFrame::Create( + params.stream, params.buffer_allocations, *allocator_, + absl::MakeSpan(args_), absl::MakeSpan(results_), + executable_->program_shape())); + + (*params.execution_scoped_state)[thunk_info().thunk_id] = + std::move(call_frame); + return initialization_status; } @@ -547,19 +573,17 @@ absl::Status HostExecuteStartThunk::ExecuteOnStream( async_events_->CreateEvent(params.host_to_device_stream->parent(), RunId(params.execution_id))); - ABSL_ASSIGN_OR_RETURN( - auto tmp_call_frame, - HostExecuteCallFrame::Create( - params.device_to_host_stream, params.host_to_device_stream, - params.buffer_allocations, *allocator_, absl::MakeSpan(args_), - absl::MakeSpan(results_), executable_->program_shape())); - - // We are making a shared pointer here because `execute` needs to be - // copyable so that it can be scheduled on the thread pool. - auto call_frame = - std::make_shared(std::move(tmp_call_frame)); - - auto execute = [this, call_frame = std::move(call_frame), params, + auto it = params.execution_scoped_state->find(thunk_info().thunk_id); + if (it == params.execution_scoped_state->end()) { + return absl::InternalError("Unable to get HostExecutableCallFrame"); + } + std::shared_ptr& call_frame = + *tsl::any_cast>(&it->second); + + ABSL_RETURN_IF_ERROR(call_frame->CopyArguments( + device_to_host_stream, params.buffer_allocations, absl::MakeSpan(args_))); + + auto execute = [this, call_frame, params, // We skip reference counting because destroying the event // would trigger a CUDA API call which is not allowed in host // callbacks. @@ -585,7 +609,8 @@ absl::Status HostExecuteStartThunk::ExecuteOnStream( return; } } - auto publish_result_status = std::move(*call_frame).PublishResult(); + auto publish_result_status = + call_frame->PublishResult(call_frame, params.buffer_allocations); if (!publish_result_status.ok()) { execute_event_ptr.SetError(publish_result_status); return; diff --git a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc index 971a9462fddb11..b402833e8244a8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include #include "absl/base/casts.h" #include "absl/container/inlined_vector.h" +#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/status_matchers.h" @@ -53,9 +54,10 @@ limitations under the License. #include "xla/stream_executor/platform_manager.h" #include "xla/stream_executor/stream_executor.h" #include "xla/stream_executor/stream_executor_address_allocator.h" +#include "xla/tests/hlo_test_base.h" #include "xla/tests/literal_test_util.h" #include "xla/tsl/concurrency/async_value_ref.h" -#include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/test_benchmark.h" #include "xla/tsl/util/proto/proto_matchers.h" #include "xla/util.h" @@ -119,8 +121,8 @@ TEST(HostExecuteStartThunkTest, SingleArgSingleResult) { se::DeviceAddressBase arg = stream_executor->Allocate(1 * sizeof(int32_t)); se::DeviceAddressBase result = stream_executor->Allocate(1 * sizeof(int32_t)); - TF_ASSERT_OK(stream->Memset32(&arg, 5, 4)); - TF_ASSERT_OK(stream->MemZero(&result, 4)); + ASSERT_OK(stream->Memset32(&arg, 5, 4)); + ASSERT_OK(stream->MemZero(&result, 4)); // Prepare buffer allocations for recording command buffer. BufferAllocation alloc_arg(/*index=*/0, 4, /*color=*/0); @@ -144,13 +146,17 @@ TEST(HostExecuteStartThunkTest, SingleArgSingleResult) { BufferAllocations allocations({arg, result}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( @@ -158,12 +164,12 @@ TEST(HostExecuteStartThunkTest, SingleArgSingleResult) { tsl::BlockUntilReady(execute_event); EXPECT_FALSE(execute_event.IsError()); - TF_ASSERT_OK(stream->WaitFor(execute_event.get().get())); - TF_ASSERT_OK(stream->BlockHostUntilDone()); + ASSERT_OK(stream->WaitFor(execute_event.get().get())); + ASSERT_OK(stream->BlockHostUntilDone()); xla::Literal result_literal(ShapeUtil::MakeShape(S32, {})); - TF_ASSERT_OK(stream->Memcpy(result_literal.untyped_data(), result, - ShapeUtil::ByteSizeOf(result_literal.shape()))); + ASSERT_OK(stream->Memcpy(result_literal.untyped_data(), result, + ShapeUtil::ByteSizeOf(result_literal.shape()))); EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR0(10), result_literal)); } @@ -193,10 +199,10 @@ TEST(HostExecuteStartThunkTest, MultiArgMultipleResult) { se::DeviceAddressBase result1 = stream_executor->Allocate(1 * sizeof(int32_t)); - TF_ASSERT_OK(stream->Memset32(&arg0, 5, 4)); - TF_ASSERT_OK(stream->Memset32(&arg1, 3, 4)); - TF_ASSERT_OK(stream->MemZero(&result0, 4)); - TF_ASSERT_OK(stream->MemZero(&result1, 4)); + ASSERT_OK(stream->Memset32(&arg0, 5, 4)); + ASSERT_OK(stream->Memset32(&arg1, 3, 4)); + ASSERT_OK(stream->MemZero(&result0, 4)); + ASSERT_OK(stream->MemZero(&result1, 4)); // Prepare buffer allocations for recording command buffer. BufferAllocation alloc_arg0(/*index=*/0, 4, /*color=*/0); @@ -225,13 +231,17 @@ TEST(HostExecuteStartThunkTest, MultiArgMultipleResult) { executable_run_options); BufferAllocations allocations({arg0, result0, arg1, result1}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( @@ -239,18 +249,18 @@ TEST(HostExecuteStartThunkTest, MultiArgMultipleResult) { tsl::BlockUntilReady(execute_event); EXPECT_FALSE(execute_event.IsError()); - TF_ASSERT_OK(stream->WaitFor(execute_event.get().get())); - TF_ASSERT_OK(stream->BlockHostUntilDone()); + ASSERT_OK(stream->WaitFor(execute_event.get().get())); + ASSERT_OK(stream->BlockHostUntilDone()); xla::Literal result_literal0(ShapeUtil::MakeShape(S32, {})); - TF_ASSERT_OK(stream->Memcpy(result_literal0.untyped_data(), result0, - ShapeUtil::ByteSizeOf(result_literal0.shape()))); + ASSERT_OK(stream->Memcpy(result_literal0.untyped_data(), result0, + ShapeUtil::ByteSizeOf(result_literal0.shape()))); EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR0(8), result_literal0)); xla::Literal result_literal1(ShapeUtil::MakeShape(S32, {})); - TF_ASSERT_OK(stream->Memcpy(result_literal1.untyped_data(), result1, - ShapeUtil::ByteSizeOf(result_literal1.shape()))); + ASSERT_OK(stream->Memcpy(result_literal1.untyped_data(), result1, + ShapeUtil::ByteSizeOf(result_literal1.shape()))); EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR0(15), result_literal1)); } @@ -308,21 +318,25 @@ TEST(HostExecuteStartThunkTest, ArgAndResultPinnedOnHost) { executable_run_options); BufferAllocations allocations({arg, result}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( stream_executor, RunId(params.execution_id))); tsl::BlockUntilReady(execute_event); EXPECT_FALSE(execute_event.IsError()); - TF_ASSERT_OK(stream->WaitFor(execute_event.get().get())); - TF_ASSERT_OK(stream->BlockHostUntilDone()); + ASSERT_OK(stream->WaitFor(execute_event.get().get())); + ASSERT_OK(stream->BlockHostUntilDone()); EXPECT_EQ( *static_cast(result_memory_allocation->address().opaque()), 10); @@ -383,21 +397,25 @@ TEST(HostExecuteStartThunkTest, ArgAndResultInSharedMemory) { executable_run_options); BufferAllocations allocations({arg, result}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( stream_executor, RunId(params.execution_id))); tsl::BlockUntilReady(execute_event); EXPECT_FALSE(execute_event.IsError()); - TF_ASSERT_OK(stream->WaitFor(execute_event.get().get())); - TF_ASSERT_OK(stream->BlockHostUntilDone()); + ASSERT_OK(stream->WaitFor(execute_event.get().get())); + ASSERT_OK(stream->BlockHostUntilDone()); EXPECT_EQ( *static_cast(result_memory_allocation->address().opaque()), 10); @@ -445,21 +463,25 @@ TEST(HostExecuteStartThunkTest, ArgAndResultNonRegisteredHostMemory) { executable_run_options); BufferAllocations allocations({arg, result}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( stream_executor, RunId(params.execution_id))); tsl::BlockUntilReady(execute_event); EXPECT_FALSE(execute_event.IsError()); - TF_ASSERT_OK(stream->WaitFor(execute_event.get().get())); - TF_ASSERT_OK(stream->BlockHostUntilDone()); + ASSERT_OK(stream->WaitFor(execute_event.get().get())); + ASSERT_OK(stream->BlockHostUntilDone()); EXPECT_EQ(result_value, 10); } @@ -514,13 +536,17 @@ TEST(HostExecuteStartThunkTest, TestErrorPropagationFromExecuteEvent) { executable_run_options); BufferAllocations allocations({arg, result}, 0, &allocator); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); - TF_ASSERT_OK( - thunk->Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk->ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk->Initialize(init_params)); + ASSERT_OK(thunk->ExecuteOnStream(params)); ASSERT_OK_AND_ASSIGN(auto execute_event, thunk->async_events()->ExtractEvent( @@ -544,9 +570,10 @@ TEST(HostExecuteDoneThunkTest, WaitingOnAvailableEvent) { BufferAllocations allocations({}, 0, nullptr); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); { ASSERT_OK_AND_ASSIGN( auto available_event, @@ -555,9 +582,12 @@ TEST(HostExecuteDoneThunkTest, WaitingOnAvailableEvent) { available_event.SetStateConcrete(); } - TF_ASSERT_OK( - thunk.Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); - TF_ASSERT_OK(thunk.ExecuteOnStream(params)); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk.Initialize(init_params)); + ASSERT_OK(thunk.ExecuteOnStream(params)); } TEST(HostExecuteDoneThunkTest, WaitingOnErrorEvent) { @@ -575,9 +605,10 @@ TEST(HostExecuteDoneThunkTest, WaitingOnErrorEvent) { BufferAllocations allocations({}, 0, nullptr); + Thunk::ExecutionScopedState scoped_state; Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( service_executable_run_options, allocations, stream.get(), stream.get(), - nullptr, nullptr, nullptr); + nullptr, nullptr, nullptr, {}, &scoped_state); { ASSERT_OK_AND_ASSIGN( auto error_event, @@ -585,8 +616,11 @@ TEST(HostExecuteDoneThunkTest, WaitingOnErrorEvent) { error_event.SetError(Internal("Test error")); } - TF_ASSERT_OK( - thunk.Initialize(Thunk::InitializeParams{/*executor=*/stream_executor})); + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + ASSERT_OK(thunk.Initialize(init_params)); EXPECT_THAT(thunk.ExecuteOnStream(params), absl_testing::StatusIs(absl::StatusCode::kInternal)); } @@ -736,6 +770,118 @@ TEST(HostExecuteStartThunkTest, IsHostOffloadSet) { EXPECT_EQ(it->second, "true"); } +using HostExecuteThunkExecuteTest = HloTestBase; + +TEST_F(HostExecuteThunkExecuteTest, RunInLoop) { + const absl::string_view hlo_string = R"( +HloModule host_execute_loop_example + +%host_computation (val: s32[]) -> s32[] { + %val = s32[] parameter(0) + %one = s32[] constant(1) + + ROOT %result = s32[] add(%val, %one) +} + +%while_cond (state: s32[]) -> pred[] { + %state = s32[] parameter(0) + %limit = s32[] constant(1000) + + ROOT %cond = pred[] compare(%state, %limit), direction=LT +} + +%while_body (state: s32[]) -> s32[] { + %state = s32[] parameter(0) + + %host_execute_start = ((s32[]), s32[]) custom-call-start(%state), + custom_call_target="HostExecute", + called_computations={%host_computation}, + async_execution_thread="host" + ROOT %host_execute_result = s32[] custom-call-done(%host_execute_start) +} + +ENTRY %main () -> s32[] { + %zero = s32[] constant(0) + + ROOT %final_state = s32[] while(%zero), + condition=%while_cond, + body=%while_body +} +)"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_string)); + + ASSERT_OK_AND_ASSIGN(Literal output, Execute(std::move(module), {})); + + EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR0(S32, 1000), + output)); +} + +void BM_HostExecuteThunkOverhead(benchmark::State& state) { + se::StreamExecutor* stream_executor = GpuExecutor(); + ASSERT_OK_AND_ASSIGN(auto stream, stream_executor->CreateStream()); + + static constexpr char const* kHloModule = R"( + HloModule module + ENTRY add_inplace { + p0 = s32[] parameter(0) + ROOT add = s32[] add(p0, p0) + } + )"; + + ASSERT_OK_AND_ASSIGN(auto hlo_module, + ParseAndReturnUnverifiedModule(kHloModule, {})); + + se::DeviceAddressBase arg = stream_executor->Allocate(1 * sizeof(int32_t)); + se::DeviceAddressBase result = stream_executor->Allocate(1 * sizeof(int32_t)); + + CHECK_OK(stream->Memset32(&arg, 5, 4)); + CHECK_OK(stream->MemZero(&result, 4)); + + // Prepare buffer allocations for recording command buffer. + BufferAllocation alloc_arg(/*index=*/0, 4, /*color=*/0); + BufferAllocation alloc_result(/*index=*/1, 4, /*color=*/0); + + BufferAllocation::Slice slice_arg(&alloc_arg, 0, 4); + BufferAllocation::Slice slice_result(&alloc_result, 0, 4); + + stream_executor::StreamExecutorAddressAllocator allocator(stream_executor); + ExecutableRunOptions executable_run_options; + executable_run_options.set_device_to_host_stream(stream.get()); + executable_run_options.set_host_to_device_stream(stream.get()); + ServiceExecutableRunOptions service_executable_run_options( + executable_run_options); + + BufferAllocations allocations({arg, result}, 0, &allocator); + + Thunk::ExecutionScopedState scoped_state; + Thunk::InitializeParams init_params{.executor = stream_executor, + .buffer_allocations = &allocations, + .stream = stream.get(), + .execution_scoped_state = &scoped_state}; + Thunk::ExecuteParams exec_params = Thunk::ExecuteParams::Create( + service_executable_run_options, allocations, stream.get(), stream.get(), + nullptr, nullptr, nullptr, {}, &scoped_state); + + ASSERT_OK_AND_ASSIGN(auto start_thunk, + CreateHostExecuteStartThunk( + Thunk::ThunkInfo(), *hlo_module, + {{slice_arg, ShapeUtil::MakeShape(S32, {})}}, + {{slice_result, ShapeUtil::MakeShape(S32, {})}})); + CHECK_OK(start_thunk->Initialize(init_params)); + + HostExecuteDoneThunk done_thunk(Thunk::ThunkInfo(), + start_thunk->async_events()); + CHECK_OK(done_thunk.Initialize(init_params)); + + for (auto s : state) { + CHECK_OK(start_thunk->ExecuteOnStream(exec_params)); + CHECK_OK(done_thunk.ExecuteOnStream(exec_params)); + } +} +BENCHMARK(BM_HostExecuteThunkOverhead); + } // namespace } // namespace gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 3d26b3168889e9..58c39a4071ce6d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -1979,6 +1979,7 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc index a61accc39ba232..e16ab8303037ea 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc @@ -31,6 +31,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/container/inlined_vector.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" @@ -1209,15 +1210,6 @@ bool IsBinaryElementwiseOfBroadcastParamOrConst(const HloInstruction& hlo) { return false; } -FusionDecision ShouldFuseUser(mlir::MLIRContext& mlir_context, - HloInstruction* user, - const HloInstruction& original_user, - HloInstruction* fusion) { - if (triton_fusion::IsOutputWorthFusing(original_user)) { - return CanFuse(mlir_context, fusion, user); - } - return FusionDecision::Forbid("Not obviously profitable to fuse as output."); -} // Holds shape tracking information for an instruction during backward BFS. struct TrackerInfo { @@ -1387,6 +1379,91 @@ std::optional ComputeCandidateTracker( return std::nullopt; } +FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose, + const HloInstruction& fusion, + const ShapeTracker& tracker) { + const HloInstruction* dot = hlo_query::FindInstruction( + fusion.fused_instructions_computation(), HloOpcode::kDot); + if (dot == nullptr) { + return FusionDecision::Forbid("Dot not found in fusion."); + } + + absl::StatusOr> dot_dims = + DotOperandDims::FromDot(dot); + if (!dot_dims.ok()) { + return FusionDecision::Forbid("Failed to get dot operand dims."); + } + const auto& [lhs_dims, rhs_dims] = *dot_dims; + + int64_t num_batch = lhs_dims.Rank(DotOperandDims::kBatch); + int64_t num_lhs_nc = lhs_dims.Rank(DotOperandDims::kNonContracting); + int64_t num_rhs_nc = rhs_dims.Rank(DotOperandDims::kNonContracting); + + // Dot result always maps to [batch..., lhs_nc..., rhs_nc...]. + absl::InlinedVector batch_dims(num_batch); + absl::c_iota(batch_dims, 0); + + absl::InlinedVector lhs_nc_dims(num_lhs_nc); + absl::c_iota(lhs_nc_dims, num_batch); + + absl::InlinedVector rhs_nc_dims(num_rhs_nc); + absl::c_iota(rhs_nc_dims, num_batch + num_lhs_nc); + + ShapeTracker transpose_tracker = tracker; + if (!transpose_tracker.AppendInstruction(&transpose).ok()) { + return FusionDecision::Forbid("Failed to append transpose to tracker."); + } + + for (auto& dim : {batch_dims, lhs_nc_dims, rhs_nc_dims}) { + if (!dim.empty() && + !transpose_tracker.MapsToOneStride(dim, /*allow_swaps=*/true)) { + return FusionDecision::Forbid("Dimension has non-contiguous section."); + } + } + + return FusionDecision::Allow(); +} + +// Propagates shape tracking information forward from the dot output to `user`. +// Returns the updated ShapeTracker if propagation is successful, or nullopt if +// propagation fails. +std::optional ComputeUserTracker( + const HloInstruction* user, std::optional current_tracker) { + if (!current_tracker.has_value()) { + return std::nullopt; + } + if (current_tracker->AppendInstruction(user).ok()) { + return current_tracker; + } + if (user->IsElementwise()) { + current_tracker->SetElementType(user->shape().element_type()); + return current_tracker; + } + return std::nullopt; +} + +FusionDecision ShouldFuseUser(const HloInstruction* user, + const HloInstruction& original_user, + const HloInstruction* fusion, + const std::optional& tracker) { + switch (user->opcode()) { + case HloOpcode::kTranspose: + if (!tracker.has_value()) { + return FusionDecision::Forbid( + "No shape tracker found for transpose user."); + } + return ShouldFuseUserTranspose(*user, *fusion, *tracker); + default: + break; + } + + if (!triton_fusion::IsOutputWorthFusing(original_user)) { + return FusionDecision::Forbid( + "Not obviously profitable to fuse as output."); + } + return FusionDecision::Allow(); +} + // Attempts to fuse all candidates and their operands into the fusion. absl::Status FuseOperandsBFS( mlir::MLIRContext& mlir_context, FusionSearchSpace& search_space, @@ -1524,6 +1601,9 @@ absl::StatusOr> CreateTileableFusion( ABSL_RETURN_IF_ERROR( FuseOperandsBFS(mlir_context, fusion_search_space, queue, fusion)); + // Initialize tracker for dot users. + std::optional epilogue_tracker = ShapeTracker(dot->shape()); + // Fuse in users until we cannot tile or reach the root. while (!fusion->IsRoot()) { // Search space was created so that the result only ever has a single user. @@ -1543,12 +1623,14 @@ absl::StatusOr> CreateTileableFusion( HloInstruction* original_user = fusion_search_space.fused_to_original().at(user); if (FusionDecision decision = - ShouldFuseUser(mlir_context, user, *original_user, fusion); + ShouldFuseUser(user, *original_user, fusion, epilogue_tracker) + .And(CanFuse(mlir_context, fusion, user)); !decision.IsAllowed()) { VLOG(5) << "Not fusing user: " << decision.Explain(); break; } VLOG(5) << "Fusing user into epilogue: " << user->ToString(); + epilogue_tracker = ComputeUserTracker(user, std::move(epilogue_tracker)); ABSL_ASSIGN_OR_RETURN( fusion, FuseUserAndOperands(mlir_context, fusion_search_space, fusion, user)); diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc index 5113eefcd9eb41..c1c6e89ad2b2bc 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc @@ -60,6 +60,12 @@ auto TransposeOrBitcastTranspose() { m::Bitcast(m::Transpose(m::Parameter()))); } +template +auto TransposeOrTransposeBitcast(Pattern pattern) { + return m::AnyOf(m::Transpose(pattern), + m::Transpose(m::Bitcast(pattern))); +} + class GemmFusionTestBase : public HloHardwareIndependentTestBase { public: GemmFusionTestBase() @@ -2450,6 +2456,93 @@ ENTRY main { GmockMatch(TransposeOrBitcastTranspose()))); } +TEST_P(GemmFusionProfitabilityTest, + DisallowEpilogueTransposeSplittingRhsNcDims) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +ENTRY e { + p0 = f16[8,64,32]{2,1,0} parameter(0) + c0 = f32[8,64,32]{2,1,0} convert(p0) + p1 = f32[8,64,16]{2,1,0} parameter(1) + dot = f32[8,32,16]{2,1,0} dot(c0, p1), + lhs_batch_dims={0}, lhs_contracting_dims={1}, + rhs_batch_dims={0}, rhs_contracting_dims={1} + bitcast = f32[8,32,2,8]{3,2,1,0} bitcast(dot) + // Transpose splits RHS non-contracting dimension N (dims 2,3) by interleaving + // batch dimension B (dim 0) between them as [N0, B, N1, M]. + ROOT transpose = f32[2,8,8,32]{3,2,1,0} transpose(bitcast), + dimensions={2,0,3,1} +})")); + ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(TransposeOrTransposeBitcast(m::Fusion()))); +} + +TEST_P(GemmFusionProfitabilityTest, + DisallowEpilogueTransposeSplittingBatchDims) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +ENTRY e { + p0 = f16[8,64,32]{2,1,0} parameter(0) + c0 = f32[8,64,32]{2,1,0} convert(p0) + p1 = f32[8,64,16]{2,1,0} parameter(1) + dot = f32[8,32,16]{2,1,0} dot(c0, p1), + lhs_batch_dims={0}, lhs_contracting_dims={1}, + rhs_batch_dims={0}, rhs_contracting_dims={1} + bitcast = f32[2,4,32,16]{3,2,1,0} bitcast(dot) + // Transpose splits batch dimension B (dims 0,1) by interleaving LHS + // non-contracting dimension M (dim 2) between them as [B0, M, B1, N]. + ROOT transpose = f32[2,32,4,16]{3,2,1,0} transpose(bitcast), + dimensions={0,2,1,3} +})")); + ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(TransposeOrTransposeBitcast(m::Fusion()))); +} + +TEST_P(GemmFusionProfitabilityTest, + DisallowEpilogueTransposeSplittingLhsNcDims) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +ENTRY e { + p0 = f16[32,64]{1,0} parameter(0) + c0 = f32[32,64]{1,0} convert(p0) + p1 = f32[16,64]{1,0} parameter(1) + dot = f32[32,16]{1,0} dot(c0, p1), + lhs_contracting_dims={1}, rhs_contracting_dims={1} + bitcast = f32[4,8,16]{2,1,0} bitcast(dot) + // Transpose splits LHS non-contracting dimension M (dims 0,1) by interleaving + // RHS non-contracting dimension N (dim 2) between them as [M0, N, M1]. + ROOT transpose = f32[4,16,8]{2,1,0} transpose(bitcast), dimensions={0,2,1} +})")); + ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(TransposeOrTransposeBitcast(m::Fusion()))); +} + +TEST_P(GemmFusionTestV2, AllowEpilogueTransposeWithSwapsWithinDimensionGroup) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +ENTRY e { + p0 = f32[32,64]{1,0} parameter(0) + p1 = f32[64,64]{1,0} parameter(1) + dot = f32[32,64]{1,0} dot(p0, p1), + lhs_contracting_dims={1}, rhs_contracting_dims={1} + bitcast = f32[1,1,32,16,4]{4,3,2,1,0} bitcast(dot) + // Transposing {0,1,4,3,2} swaps non-contracting groups (M, N) -> (N, M) + // while keeping each group internally contiguous, which is permitted. + ROOT transpose = f32[1,1,4,16,32]{4,3,2,1,0} transpose(bitcast), + dimensions={0,1,4,3,2} +})")); + ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(m::Fusion())); +} + TEST_P(GemmFusionTestV2, ConcatResetTrackerCrash) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( HloModule m diff --git a/third_party/xla/xla/literal_util.cc b/third_party/xla/xla/literal_util.cc index 54137a9122e174..84eb7d92596081 100644 --- a/third_party/xla/xla/literal_util.cc +++ b/third_party/xla/xla/literal_util.cc @@ -541,6 +541,87 @@ using RngT = std::conditional_t< sizeof(IntT) < sizeof(uint16_t), std::conditional_t::is_signed, int16_t, uint16_t>, IntT>; + +// Computes safe [min, max] bounds for integral literal generation. +// - If no limit is specified, returns the full type range [lowest, max]. +// - If use_large_range is true or bit_width <= 4, clamps the limit directly. +// - Otherwise, bounds values to B/2 bits to prevent hardware ALU overflow +// (multiplication, squaring, etc.) and float conversion explosions. +// - Expands to full range if no_duplicates requires more unique elements than +// B/2 capacity. +template +std::pair GetIntegralBounds( + const Shape& shape, bool use_large_range, bool no_duplicates, + std::optional> limit) { + if (!limit.has_value()) { + return {std::numeric_limits::lowest(), + std::numeric_limits::max()}; + } + + constexpr int64_t bit_width = sizeof(IntT) * 8; + + // Sub-byte integers (<= 4 bits) already have tiny domains (<= 16 values). + if (use_large_range || bit_width <= 4) { + return {SafeClampInt64(limit->first), + SafeClampInt64(limit->second)}; + } + + // Calculate default B/2 bitwidth bounds and default_range_size (2^H - 1), + // which is the width of the domain [default_min, default_max]. + int64_t h = bit_width / 2; + int64_t default_min; + int64_t default_max; + int64_t default_range_size; + if constexpr (std::numeric_limits::is_signed) { + default_min = -(int64_t{1} << (h - 1)); + default_max = (int64_t{1} << (h - 1)) - 1; + default_range_size = (int64_t{1} << h) - 1; + } else { + default_min = 0; + default_max = (int64_t{1} << h) - 1; + default_range_size = default_max; + } + + // If no_duplicates is requested, ensure capacity >= element count. + int64_t num_elements = ShapeUtil::ElementsIn(shape); + if (no_duplicates && num_elements > default_range_size) { + if (limit.has_value()) { + return {SafeClampInt64(limit->first), + SafeClampInt64(limit->second)}; + } + return {std::numeric_limits::lowest(), + std::numeric_limits::max()}; + } + + int64_t min_64 = default_min; + int64_t max_64 = default_max; + + if (limit.has_value()) { + bool lower_unconstrained = + (limit->first == std::numeric_limits::min()); + bool upper_unconstrained = + (limit->second == std::numeric_limits::max()); + + if (lower_unconstrained && upper_unconstrained) { + min_64 = default_min; + max_64 = default_max; + } else if (lower_unconstrained) { + max_64 = limit->second; + min_64 = + (max_64 < default_min) ? max_64 - default_range_size : default_min; + } else if (upper_unconstrained) { + min_64 = limit->first; + max_64 = + (min_64 > default_max) ? min_64 + default_range_size : default_max; + } else { + min_64 = limit->first; + max_64 = limit->second; + } + } + + return {SafeClampInt64(min_64), SafeClampInt64(max_64)}; +} + template void PopulateWithRandomIntegralDataWithBounds( Literal* literal, std::minstd_rand0* engine, bool no_duplicates, IntT min, @@ -943,12 +1024,8 @@ absl::StatusOr MakeFakeLiteral( } if constexpr (primitive_util::IsIntegralType( primitive_type_constant)) { - NativeT max = std::numeric_limits::max(); - NativeT min = std::numeric_limits::lowest(); - if (limit.has_value()) { - min = SafeClampInt64(limit->first); - max = SafeClampInt64(limit->second); - } + auto [min, max] = GetIntegralBounds( + new_shape, use_large_range, no_duplicates, limit); if (max_bits_of_precision.has_value()) { max = std::min(max, static_cast(1 << *max_bits_of_precision)); diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc index de406e695faaec..93e130039c020d 100644 --- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc +++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc @@ -27,6 +27,7 @@ limitations under the License. #include #include +#include "absl/base/attributes.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" @@ -2579,6 +2580,17 @@ PJRT_Error* PJRT_Executable_GetCompiledMemoryStats( return nullptr; } +ABSL_ATTRIBUTE_NOINLINE +absl::StatusOr>> +ParseOptionalCompileOptions(absl::string_view options_str) { + if (options_str.empty() || !options_str.data()) { + return std::make_unique>(std::nullopt); + } + ABSL_ASSIGN_OR_RETURN(auto options, ParseCompileOptions(options_str)); + return std::make_unique>( + std::move(options)); +} + PJRT_Error* PJRT_Executable_DeserializeAndLoad( PJRT_Executable_DeserializeAndLoad_Args* args) { // TODO: b/516902012 - Make this check stricter after 12week compatibility @@ -2591,16 +2603,11 @@ PJRT_Error* PJRT_Executable_DeserializeAndLoad( absl::string_view serialized(args->serialized_executable, args->serialized_executable_size); - std::optional overridden_options; - - if (args->overridden_serialized_compile_options && - args->overridden_serialized_compile_options_size > 0) { - PJRT_ASSIGN_OR_RETURN( - overridden_options, - ParseCompileOptions(absl::string_view( - args->overridden_serialized_compile_options, - args->overridden_serialized_compile_options_size))); - } + PJRT_ASSIGN_OR_RETURN( + std::unique_ptr> overridden_options, + ParseOptionalCompileOptions( + absl::string_view(args->overridden_serialized_compile_options, + args->overridden_serialized_compile_options_size))); xla::LoadOptions load_options; if (args->struct_size >= @@ -2623,9 +2630,10 @@ PJRT_Error* PJRT_Executable_DeserializeAndLoad( } } - PJRT_ASSIGN_OR_RETURN(std::unique_ptr executable, - args->client->client->LoadSerializedExecutable( - serialized, overridden_options, load_options)); + PJRT_ASSIGN_OR_RETURN( + std::unique_ptr executable, + args->client->client->LoadSerializedExecutable( + serialized, std::move(*overridden_options), load_options)); args->loaded_executable = new PJRT_LoadedExecutable(std::move(executable), args->client); diff --git a/third_party/xla/xla/pjrt/pjrt_client_test.cc b/third_party/xla/xla/pjrt/pjrt_client_test.cc index 6c466c32467ea7..7f6b62f30afb13 100644 --- a/third_party/xla/xla/pjrt/pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/pjrt_client_test.cc @@ -830,17 +830,21 @@ TEST(PjRtClientTest, ClearPeakMemory) { ASSERT_OK(buffer->GetReadyFuture().Await()); TF_ASSERT_OK_AND_ASSIGN(auto alloc_stats, device->GetAllocatorStats()); - EXPECT_EQ(alloc_stats.bytes_in_use, initial_stats.bytes_in_use + kAllocSize); - EXPECT_EQ(alloc_stats.peak_bytes_in_use, + ASSERT_EQ(alloc_stats.bytes_in_use, initial_stats.bytes_in_use + kAllocSize); + ASSERT_EQ(alloc_stats.peak_bytes_in_use, initial_stats.peak_bytes_in_use + kAllocSize); - EXPECT_EQ(alloc_stats.bytes_in_use, alloc_stats.peak_bytes_in_use); + ASSERT_EQ(alloc_stats.bytes_in_use, alloc_stats.peak_bytes_in_use); + ASSERT_EQ(alloc_stats.peak_allocated_bytes, + initial_stats.peak_allocated_bytes + kAllocSize); // dealloc buffer.reset(); TF_ASSERT_OK_AND_ASSIGN(auto dealloc_stats, device->GetAllocatorStats()); - EXPECT_EQ(initial_stats.bytes_in_use, dealloc_stats.bytes_in_use); - EXPECT_EQ(dealloc_stats.peak_bytes_in_use, alloc_stats.peak_bytes_in_use); + ASSERT_EQ(initial_stats.bytes_in_use, dealloc_stats.bytes_in_use); + ASSERT_EQ(dealloc_stats.peak_bytes_in_use, alloc_stats.peak_bytes_in_use); + ASSERT_EQ(dealloc_stats.peak_allocated_bytes, + alloc_stats.peak_allocated_bytes); absl::Status clear_status = device->ClearMemoryStats(); if (!absl::IsUnimplemented(clear_status)) { @@ -848,6 +852,8 @@ TEST(PjRtClientTest, ClearPeakMemory) { TF_ASSERT_OK_AND_ASSIGN(auto clear_stats, device->GetAllocatorStats()); EXPECT_EQ(clear_stats.bytes_in_use, dealloc_stats.bytes_in_use); EXPECT_EQ(clear_stats.peak_bytes_in_use, dealloc_stats.bytes_in_use); + EXPECT_EQ(clear_stats.peak_allocated_bytes, + dealloc_stats.bytes_in_use + dealloc_stats.bytes_reserved); } } struct LinearizePackTestParam { diff --git a/third_party/xla/xla/python/ifrt/ir/BUILD b/third_party/xla/xla/python/ifrt/ir/BUILD index a5724060b3bd01..7b7d1a3d734e7d 100644 --- a/third_party/xla/xla/python/ifrt/ir/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/BUILD @@ -757,6 +757,7 @@ cc_library( ":ir", "//xla:status_macros", "//xla/pjrt:host_memory_spaces", + "//xla/pjrt:pjrt_compiler", "//xla/python/ifrt", "//xla/python/ifrt:attribute_map", "//xla/python/ifrt:remap_plan_proto_cc", diff --git a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc index d65d79ac9d3c7f..944cf95a33d54a 100644 --- a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc +++ b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc @@ -46,6 +46,7 @@ limitations under the License. #include "mlir/Support/DebugStringHelper.h" #include "mlir/Support/LLVM.h" #include "xla/pjrt/host_memory_spaces.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/python/ifrt/array.h" #include "xla/python/ifrt/array_spec.h" #include "xla/python/ifrt/attribute_map.h" @@ -376,8 +377,8 @@ struct CallLoadedExecutableOpState { std::vector arrays_to_remove; { - std::vector non_donatable_pinned_host_inputs; - std::vector non_donatable_pinned_host_inputs_handles; + std::vector to_copy_non_donatable_inputs; + std::vector to_copy_non_donatable_inputs_handles; for (int idx = 0; idx < input_handles.size(); ++idx) { const ArrayHandle handle = input_handles[idx]; @@ -399,12 +400,17 @@ struct CallLoadedExecutableOpState { "Input will not be donated. \n" << pretty_print; // TODO(b/401105456): Do not special case pinned host arrays once - // non-donatable pinned host inputs are supported. + // non-donatable pinned host inputs are supported on non-CPU devices. if (!is_mpmd_reshard && array_it->second.array->sharding().memory_kind() == - kPinnedHostMemoryKind) { - non_donatable_pinned_host_inputs.push_back(array_it->second.array); - non_donatable_pinned_host_inputs_handles.push_back(handle); + kPinnedHostMemoryKind && + array_it->second.array->sharding() + .devices() + ->devices() + .front() + ->PlatformName() != xla::CpuName()) { + to_copy_non_donatable_inputs.push_back(array_it->second.array); + to_copy_non_donatable_inputs_handles.push_back(handle); } else { options.non_donatable_input_indices.insert(idx); } @@ -420,15 +426,15 @@ struct CallLoadedExecutableOpState { // TODO(b/401105456): Remove this CopyArrays call once non-donatable // pinned host inputs are supported. - if (!non_donatable_pinned_host_inputs.empty()) { + if (!to_copy_non_donatable_inputs.empty()) { ABSL_ASSIGN_OR_RETURN( std::vector copied_pinned_host_inputs, - env.client->CopyArrays( - absl::MakeSpan(non_donatable_pinned_host_inputs), - /*devices=*/std::nullopt, - /*memory_kind=*/std::nullopt, ArrayCopySemantics::kAlwaysCopy)); + env.client->CopyArrays(absl::MakeSpan(to_copy_non_donatable_inputs), + /*devices=*/std::nullopt, + /*memory_kind=*/std::nullopt, + ArrayCopySemantics::kAlwaysCopy)); for (int idx = 0; idx < copied_pinned_host_inputs.size(); ++idx) { - env.handle_to_array[non_donatable_pinned_host_inputs_handles[idx]] = + env.handle_to_array[to_copy_non_donatable_inputs_handles[idx]] = ArrayState{ /*array=*/std::move(copied_pinned_host_inputs[idx]), /*can_be_donated=*/false, diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc index 2d98119ad4ad89..1dc6a8d6787c6f 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc @@ -159,6 +159,9 @@ absl::StatusOr>> GetHloShardings( ABSL_ASSIGN_OR_RETURN( auto hlo_sharding, xla::HloSharding::FromProto((*pjrt_executable_op_shardings)[i])); + if (hlo_sharding.UseNamedShardingLeaf()) { + hlo_sharding = xla::HloSharding::V3ToV2Sharding(hlo_sharding); + } hlo_shardings.push_back(hlo_sharding); } } diff --git a/third_party/xla/xla/service/compiler.h b/third_party/xla/xla/service/compiler.h index 82a06930ee57ea..b7d37c0aac679f 100644 --- a/third_party/xla/xla/service/compiler.h +++ b/third_party/xla/xla/service/compiler.h @@ -438,21 +438,6 @@ class AotCompilationOptions { run_backend_only_ = run_backend_only; } - bool sanitize_memory() const { return sanitize_memory_; } - void set_sanitize_memory(bool sanitize_memory) { - sanitize_memory_ = sanitize_memory; - } - - int sanitize_memory_track_origins() const { - return sanitize_memory_track_origins_; - } - void set_sanitize_memory_track_origins(int track_origins) { - sanitize_memory_track_origins_ = track_origins; - if (track_origins > 0) { - sanitize_memory_ = true; - } - } - bool sanitize_dataflow() const { return sanitize_dataflow_; } void set_sanitize_dataflow(bool sanitize_dataflow) { sanitize_dataflow_ = sanitize_dataflow; @@ -500,8 +485,6 @@ class AotCompilationOptions { int64_t profile_version_ = 0; std::string cache_key_; bool run_backend_only_ = false; - bool sanitize_memory_ = false; - int sanitize_memory_track_origins_ = 0; bool sanitize_dataflow_ = false; std::vector sanitize_abilists_dataflow_; // Contains target-specific information required by AOT compilation. diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 0d0410d6ce45ee..99dc48e2494892 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -983,7 +983,6 @@ cc_library( "//xla/backends/cpu/collectives:cpu_cliques", "//xla/backends/cpu/collectives:cpu_collectives", "//xla/backends/cpu/collectives:in_process_collectives", - "//xla/backends/cpu/runtime:msan_emulated_tls", "//xla/backends/cpu/runtime:xfeed_manager", "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index 3fd1bcb5363d62..39a6808562aa7c 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -2272,9 +2272,6 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr hlo_module, IrCompiler::GetCodeGenOptLevel(hlo_module->config()); llvm::TargetOptions target_options = CompilerTargetOptions(hlo_module->config()); - if (options.sanitize_memory()) { - target_options.EmulatedTLS = true; - } auto target_machine_builder = [&]() { return absl::WrapUnique(target->createTargetMachine( triple, options.cpu_name(), options.features(), target_options, @@ -2351,8 +2348,6 @@ CpuCompiler::CompileAheadOfTimeThunks( options::DisableLoopUnrolling(module->config()), /*disable_platform_dependent_math=*/ options::DisablePlatformDependentMath(module->config()) || fast_compile, - /*msan_enabled=*/aot_options.sanitize_memory(), - /*msan_track_origins=*/aot_options.sanitize_memory_track_origins(), /*dfsan_enabled=*/aot_options.sanitize_dataflow(), /*dfsan_abilists_enabled=*/aot_options.sanitize_abilists_dataflow()}; diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h index e96883b007eeca..c2b34ff178dc79 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime.h +++ b/third_party/xla/xla/service/cpu/cpu_runtime.h @@ -38,9 +38,8 @@ namespace runtime { // Names of runtime functions. These get resolved from the generated code to the // right symbol at link time in one of two ways: -// 1. When using the JIT, the symbol resolver -// (xla::cpu::BuiltinDefinitionGenerator) maps this symbol name to the actual -// symbol. +// 1. When using the JIT, the symbol resolver (xla::cpu::RuntimeSymbolGenerator) +// maps this symbol name to the actual symbol. // 2. When using ahead-of-time compilation, the linker can resolve the name // because it is a symbol in the cpu_runtime library. inline constexpr absl::string_view kEigenMatMulF16SymbolName = @@ -139,8 +138,6 @@ inline constexpr absl::string_view kReduceScatterSymbolName = "__xla_cpu_runtime_ReduceScatter"; inline constexpr absl::string_view kHandleFfiCallSymbolName = "__xla_cpu_runtime_HandleFfiCall"; -inline constexpr absl::string_view kMsanEmutlsGetAddressBridgeSymbolName = - "__xla_cpu_runtime_emutls_get_address"; // All symbol names for XLA CPU runtime functions need to start with this // prefix. diff --git a/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc b/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc index 5ab9ee1aa374bf..aef57cadbb6b4f 100644 --- a/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc +++ b/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc @@ -18,9 +18,6 @@ limitations under the License. #include #include "absl/base/casts.h" -#ifdef ABSL_HAVE_MEMORY_SANITIZER -#include -#endif #include "absl/strings/match.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" @@ -96,10 +93,6 @@ ENTRY e { /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::BigPic); aot_options->set_executor(stream_exec); -#ifdef ABSL_HAVE_MEMORY_SANITIZER - aot_options->set_sanitize_memory(true); - aot_options->set_sanitize_memory_track_origins(__msan_get_track_origins()); -#endif auto test = [this, &compiler, aot_options = std::move(aot_options)]( absl::string_view test_name, absl::string_view hlo, int input, @@ -113,9 +106,6 @@ ENTRY e { TF_ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); -#ifdef ABSL_HAVE_MEMORY_SANITIZER - EXPECT_TRUE(absl::StrContains(serialized_aot_result, "__msan_")); -#endif TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, compiler->LoadAotCompilationResult(serialized_aot_result)); diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index bf967c98dc36cb..a15a040a367045 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1876,7 +1876,6 @@ cc_library( "//xla/backends/gpu/transforms:sanitize_constant_names", "//xla/backends/gpu/transforms:scalar_constant_sinker", "//xla/backends/gpu/transforms:scaled_dot_rewriter", - "//xla/backends/gpu/transforms:scan_rewriter", "//xla/backends/gpu/transforms:scatter_determinism_expander", "//xla/backends/gpu/transforms:scatter_expander", "//xla/backends/gpu/transforms:scatter_slice_simplifier", diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index cf971cf22f202e..3f12c4b13efb2d 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -140,7 +140,6 @@ limitations under the License. #include "xla/backends/gpu/transforms/sanitize_constant_names.h" #include "xla/backends/gpu/transforms/scalar_constant_sinker.h" #include "xla/backends/gpu/transforms/scaled_dot_rewriter.h" -#include "xla/backends/gpu/transforms/scan_rewriter.h" #include "xla/backends/gpu/transforms/scatter_determinism_expander.h" #include "xla/backends/gpu/transforms/scatter_expander.h" #include "xla/backends/gpu/transforms/scatter_slice_simplifier.h" @@ -906,13 +905,6 @@ absl::Status RunOptimizationPasses( pipeline.AddPass(); pipeline.AddPass(); - // Rewrite eligible scans to CUB device scans only after SPMD partitioning, - // so sharded scans are partitioned as scans (the partitioner replicates - // unknown custom calls, and the CUB call's tuple result crashes the Shardy - // sharding import). The scans that remain fall through to - // AssociativeScanRewriter and ScanExpander below. - pipeline.AddPass(); - int64_t rw_length = debug_options.xla_reduce_window_rewrite_base_length(); pipeline.AddPass>(rw_length); if (rw_length != 0) { diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index daa45dc3bf418e..ef2e1bfcbf9812 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// LINT: LEGACY_NAMES + // This proto file defines messages which represent the HLO module. This is a // full fidelity serialization of the c++ HLO constructs. // @@ -950,6 +952,14 @@ message HloPassMetadata { repeated KeyValueMetric kv_metrics = 11; } +message X64ConfigProto { + // Only to support logging x64 values. Not needed for original value tracking + // route and not mirrored in DebugAttributes in hlo_module.h. + string xprecision_64bit_fragment = 1; + string original_log_instruction_name = 2; + string xprecision_64bit_original_type = 3; +} + message DebugAttributesProto { enum DebugLogModeProto { // No logging is performed. @@ -959,7 +969,7 @@ message DebugAttributesProto { // not perturb the computation graph and the data logged is exactly what it // would be without logging. DEFAULT = 1; - // Same with DEFAULT, but fused value can logged at the expense of + // Same with DEFAULT, but fused value can be logged at the expense of // significant performance overhead. FUSION_DEBUGGER = 2; // Uses an optimization barrier to prevent fusion, preserving ordering but @@ -983,10 +993,9 @@ message DebugAttributesProto { // The sharding of the operands of the custom call. This is used by the // runtime to unshard tensors logged in guaranteed mode. string operands_sharding = 5; -} - -message DebugLogBackendConfigProto { - DebugAttributesProto debug_attributes_config = 64; + // The x64 config for the tensor. This is used by the runtime to reconstruct + // the 64-bit original value of the tensor in guaranteed mode. + X64ConfigProto x64_config = 6; } message DebugAttributeTableEntryProto { diff --git a/third_party/xla/xla/stream_executor/allocator_stats.cc b/third_party/xla/xla/stream_executor/allocator_stats.cc index 1d10eba776da7b..51ae5fbac915ed 100644 --- a/third_party/xla/xla/stream_executor/allocator_stats.cc +++ b/third_party/xla/xla/stream_executor/allocator_stats.cc @@ -31,6 +31,7 @@ std::string AllocatorStats::DebugString() const { "MaxAllocSize: %20s\n" "Reserved: %20s\n" "PeakReserved: %20s\n" + "PeakAllocated: %20s\n" "LargestFreeBlock: %20s\n", tsl::strings::HumanReadableNumBytes(this->bytes_limit ? *this->bytes_limit : 0), @@ -40,6 +41,7 @@ std::string AllocatorStats::DebugString() const { tsl::strings::HumanReadableNumBytes(this->largest_alloc_size), tsl::strings::HumanReadableNumBytes(this->bytes_reserved), tsl::strings::HumanReadableNumBytes(this->peak_bytes_reserved), + tsl::strings::HumanReadableNumBytes(this->peak_allocated_bytes), tsl::strings::HumanReadableNumBytes(this->largest_free_block_bytes)); } diff --git a/third_party/xla/xla/stream_executor/allocator_stats.h b/third_party/xla/xla/stream_executor/allocator_stats.h index c6d185cbcd777b..258c19b2491529 100644 --- a/third_party/xla/xla/stream_executor/allocator_stats.h +++ b/third_party/xla/xla/stream_executor/allocator_stats.h @@ -40,6 +40,7 @@ struct AllocatorStats { int64_t bytes_reserved; // Number of bytes reserved on the stack. int64_t peak_bytes_reserved; // The peak number of bytes reserved on the stack. + int64_t peak_allocated_bytes; // Peak of reserved and in-use bytes. // The upper limit on the number bytes of reservable memory on the stack, // if such a limit is known. std::optional bytes_reservable_limit; @@ -53,6 +54,7 @@ struct AllocatorStats { largest_alloc_size(0), bytes_reserved(0), peak_bytes_reserved(0), + peak_allocated_bytes(0), largest_free_block_bytes(0) {} std::string DebugString() const; diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc b/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc index 574f7ae8566039..f8779d7dc11812 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc @@ -383,6 +383,9 @@ void* GpuCudaMallocAsyncAllocator::AllocateRaw(size_t alignment, } stats_->peak_bytes_in_use = std::max(stats_->peak_bytes_in_use, stats_->bytes_in_use); + stats_->peak_allocated_bytes = + std::max(stats_->peak_allocated_bytes, + stats_->bytes_in_use + stats_->bytes_reserved); stats_->largest_alloc_size = std::max(stats_->largest_alloc_size, num_bytes); bool ptr_inserted = size_map_.emplace(ptr, num_bytes).second; @@ -463,6 +466,7 @@ bool GpuCudaMallocAsyncAllocator::ClearStats() { absl::MutexLock l(mutex_); stats_->num_allocs = 0; stats_->peak_bytes_in_use = stats_->bytes_in_use; + stats_->peak_allocated_bytes = stats_->bytes_in_use + stats_->bytes_reserved; stats_->largest_alloc_size = 0; return true; } diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 8a87fbe07d9240..3db05b09e4aed0 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -1192,6 +1192,38 @@ ENTRY main { EXPECT_LE(p0_int.max, 150.0); } +TEST_F(ConstraintPropagatorTest, GuardedOffsetLog) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + param_0 = pred[8,128] parameter(0) + param_1 = s32[8,128] parameter(1) + param_2 = s32[8,128] parameter(2) + add_s32 = s32[8,128] add(param_1, param_2) + c_zero_s32 = s32[] constant(0) + b_zero_s32 = s32[8,128] broadcast(c_zero_s32), dimensions={} + select = s32[8,128] select(param_0, add_s32, b_zero_s32) + c_neg_one = s32[] constant(-1) + b_neg_one = s32[8,128] broadcast(c_neg_one), dimensions={} + sub = s32[8,128] add(select, b_neg_one) + conv = f32[8,128] convert(sub) + c_offset = f32[] constant(1024) + b_offset = f32[8,128] broadcast(c_offset), dimensions={} + add_f32 = f32[8,128] add(b_offset, conv) + ROOT log = f32[8,128] log(add_f32) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto p1_int = states[module->entry_computation()->parameter_instruction(1)] + .GetConstraintInterval(); + auto p2_int = states[module->entry_computation()->parameter_instruction(2)] + .GetConstraintInterval(); + EXPECT_TRUE(p1_int.IsPositive()); + EXPECT_TRUE(p2_int.IsPositive()); +} + TEST_F(ConstraintPropagatorTest, MaxAddReductionElementsPerExpTracksDownstreamReductions) { const char* hlo = R"( diff --git a/third_party/xla/xla/tools/hlo_isolation/BUILD b/third_party/xla/xla/tools/hlo_isolation/BUILD index 003ef54995801b..0e32f101e730f6 100644 --- a/third_party/xla/xla/tools/hlo_isolation/BUILD +++ b/third_party/xla/xla/tools/hlo_isolation/BUILD @@ -54,7 +54,6 @@ cc_library( "//xla/tools/hlo_dump:hlo_dump_utils", "//xla/tsl/platform:env", "//xla/tsl/platform:test", - "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -102,6 +101,7 @@ xla_cc_test( "//xla:shape_util", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", + "//xla/pjrt:pjrt_executable", "//xla/service:device_assignment", "//xla/service:hlo_runner_interface", "//xla/tests:hlo_interpreter_reference_mixin", @@ -114,14 +114,14 @@ xla_cc_test( "//xla/tsl/platform:env", "//xla/tsl/platform:test", "@com_google_absl//absl/base:nullability", - "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@tsl//tsl/platform:path", diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc index 4ea25660a589fe..06dcdf06a95162 100644 --- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc +++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc @@ -29,7 +29,6 @@ limitations under the License. #include #include -#include "absl/cleanup/cleanup.h" #include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -240,11 +239,12 @@ absl::Status CompareOutputs(const HloModule& module, const Literal& test_output, return status; } -using GroupKey = std::pair; +} // namespace std::vector CreateDumpHloOutputCallbacks( - HloModule* module, const std::function& - eval_literal_mutator) { + HloModule* module, std::shared_ptr expected_literals, + const std::function& + eval_literal_mutator) { std::vector reference_callbacks; int64_t next_id = 1000000; for (auto* computation : module->computations()) { @@ -260,7 +260,8 @@ std::vector CreateDumpHloOutputCallbacks( hlo_cb.num_operands = 1; hlo_cb.callback = [hlo_name = std::string(instruction->name()), - module_name = std::string(module->name()), eval_literal_mutator]( + module_name = std::string(module->name()), expected_literals, + eval_literal_mutator]( int64_t replica_id, int64_t partition_id, absl::Span const> literals) { if (literals.empty() || !literals[0]) { @@ -269,20 +270,15 @@ std::vector CreateDumpHloOutputCallbacks( << hlo_name << " within fusion " << module_name; return; } - Literal mutated_literal = literals[0]->Clone(); + std::shared_ptr stored_literal = literals[0]; if (eval_literal_mutator) { + Literal mutated_literal = literals[0]->Clone(); eval_literal_mutator(hlo_name, &mutated_literal); + stored_literal = + std::make_shared(std::move(mutated_literal)); } - - auto* env = tsl::Env::Default(); - std::string filepath = GetFusionDebuggerFilePath(hlo_name); - auto status = tsl::WriteStringToFile( - env, filepath, mutated_literal.ToProto().SerializeAsString()); - if (!status.ok()) { - LOG(ERROR) - << "Failed to write literal to Sponge artifacts for op " - << hlo_name << " within fusion " << module_name << ": " - << status; + if (expected_literals != nullptr) { + (*expected_literals)[hlo_name] = stored_literal; } }; reference_callbacks.push_back(std::move(hlo_cb)); @@ -294,6 +290,7 @@ std::vector CreateDumpHloOutputCallbacks( std::vector CreateComparisonHloOutputCallbacks( HloModule* test_module_clone, const absl::flat_hash_map>& ref_groups, + std::shared_ptr expected_literals, const HloModule& original_module, const ModuleIsolationOptions& options, std::shared_ptr result_mutex, HloIsolationTestResult* test_result) { @@ -341,8 +338,9 @@ std::vector CreateComparisonHloOutputCallbacks( dynamic_cb.callback_id = hlo_id; dynamic_cb.num_operands = 1; dynamic_cb.callback = - [op_name, ref_op_name, module_name = original_module.name(), - abs_error, rel_error, result_mutex, test_result]( + [op_name, ref_op_name, expected_literals, + module_name = original_module.name(), abs_error, rel_error, + result_mutex, test_result]( int64_t replica_id, int64_t partition_id, absl::Span const> literals) { if (literals.empty() || !literals[0]) { @@ -353,46 +351,33 @@ std::vector CreateComparisonHloOutputCallbacks( return; } - std::string filepath = GetFusionDebuggerFilePath(ref_op_name); - - bool shape_matched = false; - Literal expected_literal; - xla::LiteralProto literal_proto; - std::string content; - if (tsl::ReadFileToString(tsl::Env::Default(), filepath, &content) - .ok()) { - if (literal_proto.ParseFromString(content)) { - auto expected_status = Literal::CreateFromProto(literal_proto); - if (expected_status.ok()) { - expected_literal = std::move(*expected_status); - if (ShapeUtil::Compatible(expected_literal.shape(), - literals[0]->shape())) { - shape_matched = true; - } - } + std::shared_ptr expected_literal_ptr; + if (expected_literals != nullptr) { + absl::MutexLock lock(result_mutex.get()); + auto it = expected_literals->find(ref_op_name); + if (it != expected_literals->end()) { + expected_literal_ptr = it->second; + expected_literals->erase(it); } } - if (!shape_matched) { - LOG(WARNING) - << "No reference literal matches the shape of the current " - "actual literal " - << literals[0]->shape().ToString() << " for op " << op_name - << " within fusion " << module_name; + if (expected_literal_ptr == nullptr) { + LOG(WARNING) << "No reference literal found in memory for op " + << ref_op_name << " within fusion " << module_name; return; } - if (expected_literal.shape().element_type() != - literals[0]->shape().element_type()) { - absl::StatusOr converted_literal_status = - expected_literal.Convert(literals[0]->shape().element_type()); - if (!converted_literal_status.ok()) { - LOG(ERROR) << "Failed to convert expected literal type for op " - << op_name << " within fusion " << module_name - << ": " << converted_literal_status.status(); - return; - } - expected_literal = std::move(*converted_literal_status); + Literal expected_literal; + if (ShapeUtil::Compatible(expected_literal_ptr->shape(), + literals[0]->shape())) { + expected_literal = expected_literal_ptr->Clone(); + } else { + LOG(WARNING) << "Reference literal shape " + << expected_literal_ptr->shape().ToString() + << " is not compatible with actual literal shape " + << literals[0]->shape().ToString() << " for op " + << op_name << " within fusion " << module_name; + return; } auto on_miscompare = @@ -436,7 +421,7 @@ std::vector CreateComparisonHloOutputCallbacks( ADD_FAILURE() << error_message; LOG(ERROR) << error_message; - absl::MutexLock lock(*result_mutex); + absl::MutexLock lock(result_mutex.get()); NumericCheck* numeric_check = test_result->add_numeric_checks(); numeric_check->set_name(absl::StrCat("FusionDebugger:", op_name)); numeric_check->set_expected_contains_inf_or_nan( @@ -457,45 +442,6 @@ std::vector CreateComparisonHloOutputCallbacks( return dynamic_cbs; } -} // namespace - -std::string GetFusionDebuggerDir() { - std::string outdir; - if (tsl::io::GetTestUndeclaredOutputsDir(&outdir)) { - return outdir; - } - std::string temp_file = tsl::io::GetTempFilename(""); - std::string temp_dir = std::string(tsl::io::Dirname(temp_file)); - (void)tsl::Env::Default()->DeleteFile(temp_file); - return temp_dir; -} - -std::string GetFusionDebuggerFilePath(absl::string_view op_name) { - std::string filename = - absl::StrCat("fusion-debugger-reference-", op_name, ".bin"); - std::string outdir; - if (tsl::io::GetTestUndeclaredOutputsDir(&outdir)) { - return tsl::io::JoinPath(outdir, filename); - } - return tsl::io::GetTempFilename(filename); -} - -void CleanUpAllFusionDebuggerFiles() { - tsl::Env* env = tsl::Env::Default(); - for (const std::string& path : GetLeftoverFusionDebuggerFiles()) { - (void)env->DeleteFile(path); - } -} - -std::vector GetLeftoverFusionDebuggerFiles() { - std::vector bin_files; - tsl::Env* env = tsl::Env::Default(); - std::string pattern = tsl::io::JoinPath(GetFusionDebuggerDir(), - "*fusion-debugger-reference-*.bin"); - (void)env->GetMatchingPaths(pattern, &bin_files); - return bin_files; -} - void PopulateNumericCheckMismatches( NumericCheck* numeric_check, const absl::StatusOr>& top_mismatches) { @@ -541,7 +487,7 @@ absl::StatusOr RunModule(std::unique_ptr module, std::vector reference_callbacks; if (options.use_fusion_debugger && options.hlo_output_callbacks.empty()) { reference_callbacks = CreateDumpHloOutputCallbacks( - module.get(), options.eval_literal_mutator); + module.get(), options.expected_literals, options.eval_literal_mutator); } ABSL_ASSIGN_OR_RETURN( @@ -703,10 +649,11 @@ absl::StatusOr RunIsolationTestOnModule( } } - absl::Cleanup cleanup = [] { CleanUpAllFusionDebuggerFiles(); }; + auto expected_literals = std::make_shared(); RunModuleOptions reference_opts; reference_opts.use_fusion_debugger = true; + reference_opts.expected_literals = expected_literals; absl::StatusOr debug_reference_output = run_module(std::move(debug_despecialized_module), reference_runner, input_data, reference_opts); @@ -726,12 +673,13 @@ absl::StatusOr RunIsolationTestOnModule( std::vector dynamic_cbs = CreateComparisonHloOutputCallbacks(test_module_clone.get(), ref_groups, - module, options, result_mutex, - test_result); + expected_literals, module, options, + result_mutex, test_result); RunModuleOptions retry_opts; retry_opts.hlo_output_callbacks = dynamic_cbs; retry_opts.use_fusion_debugger = true; + retry_opts.expected_literals = expected_literals; absl::StatusOr retry_test_output = run_module( std::move(test_module_clone), test_runner, input_data, retry_opts); if (!retry_test_output.ok()) { diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.h b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.h index f5a5a3e5230d6f..a803ab3eea474a 100644 --- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.h +++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.h @@ -22,13 +22,16 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/service/hlo_runner_interface.h" @@ -38,12 +41,18 @@ limitations under the License. namespace xla { namespace hlo_isolation { +using ExpectedLiteralsMap = + absl::flat_hash_map>; + +using GroupKey = std::pair; + struct RunModuleOptions { bool run_hlo_passes = false; bool use_fusion_debugger = false; absl::Span hlo_output_callbacks = {}; std::function eval_literal_mutator = nullptr; + std::shared_ptr expected_literals = nullptr; }; struct ModuleIsolationOptions { @@ -87,6 +96,19 @@ absl::StatusOr RunModule(std::unique_ptr module, absl::Span input_data, const RunModuleOptions& options = {}); +std::vector CreateDumpHloOutputCallbacks( + HloModule* module, std::shared_ptr expected_literals, + const std::function& + eval_literal_mutator = nullptr); + +std::vector CreateComparisonHloOutputCallbacks( + HloModule* test_module_clone, + const absl::flat_hash_map>& ref_groups, + std::shared_ptr expected_literals, + const HloModule& original_module, const ModuleIsolationOptions& options, + std::shared_ptr result_mutex, + HloIsolationTestResult* test_result); + void PopulateNumericCheckMismatches( NumericCheck* numeric_check, const absl::StatusOr>& top_mismatches); @@ -133,11 +155,6 @@ bool ComputationHasRng(const HloComputation* computation); bool LiteralContainsInfOrNan(const LiteralSlice& literal); bool ModuleContainsConstantInfOrNan(const HloModule& module); -std::string GetFusionDebuggerDir(); -std::string GetFusionDebuggerFilePath(absl::string_view op_name); -void CleanUpAllFusionDebuggerFiles(); -std::vector GetLeftoverFusionDebuggerFiles(); - } // namespace hlo_isolation } // namespace xla diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc index 6ec92a0d07f8bc..ecd7b78a58a240 100644 --- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc +++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc @@ -28,16 +28,16 @@ limitations under the License. #include #include #include "absl/base/nullability.h" -#include "absl/cleanup/cleanup.h" +#include "absl/container/flat_hash_map.h" #include "absl/functional/any_invocable.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" -#include "absl/status/status_matchers.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/array2d.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -46,6 +46,7 @@ limitations under the License. #include "xla/hlo/parser/hlo_parser.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/pjrt/pjrt_executable.h" #include "xla/service/device_assignment.h" #include "xla/service/hlo_runner_interface.h" #include "xla/shape_util.h" @@ -63,10 +64,6 @@ namespace xla { namespace hlo_isolation { namespace { -using ::absl_testing::StatusIs; -using ::testing::Contains; -using ::testing::IsEmpty; - class HloIsolationTest : public HloIsolationTestMixin> { }; @@ -860,6 +857,78 @@ ENTRY main { EXPECT_TRUE(test_runner.last_run_hlo_passes_); } +TEST_F( + HloIsolationTest, + TestRunIsolationTestOnModule_FusionDebuggerRetryPassWithExpectedLiterals) { + DelegatingRunner test_runner(&this->test_runner()); + DelegatingRunner reference_runner(&this->reference_runner()); + + const char* hlo_text = R"( +HloModule TestModule + +ENTRY main { + a = f32[] parameter(0) + b = f32[] parameter(1) + add = f32[] add(a, b) + ROOT mul = f32[] multiply(add, b) +} +)"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + + bool saw_reference_with_expected_literals = false; + bool saw_retry_with_expected_literals = false; + + ModuleIsolationOptions options; + options.run_module_fn = + [&](std::unique_ptr m, HloRunnerInterface* r, + absl::Span input_data, + const RunModuleOptions& run_opts) -> absl::StatusOr { + if (run_opts.use_fusion_debugger && run_opts.expected_literals != nullptr) { + if (run_opts.hlo_output_callbacks.empty()) { + saw_reference_with_expected_literals = true; + } else { + saw_retry_with_expected_literals = true; + } + } + + std::string module_name = std::string(m->name()); + ABSL_ASSIGN_OR_RETURN(Literal output, + RunModule(std::move(m), r, input_data, run_opts)); + + // Inject mismatch on main test runner run to trigger stage 1 and stage 2 + // failures + if (r == &test_runner && run_opts.hlo_output_callbacks.empty() && + !absl::StrContains(module_name, "defused")) { + *static_cast(output.untyped_data()) += 100.0f; + } + + return output; + }; + + std::vector args; + args.push_back(LiteralUtil::CreateR0(2.0f)); + args.push_back(LiteralUtil::CreateR0(3.0f)); + + ::testing::TestPartResultArray failures; + absl::StatusOr result_or; + { + ::testing::ScopedFakeTestPartResultReporter reporter( + ::testing::ScopedFakeTestPartResultReporter::INTERCEPT_ALL_THREADS, + &failures); + result_or = RunIsolationTestOnModule(*module, &test_runner, + &reference_runner, options, args); + } + + ASSERT_OK(result_or.status()); + EXPECT_EQ(result_or->state(), State::FAILURE); + EXPECT_EQ(result_or->reason(), "NUMERIC_MISMATCH"); + + EXPECT_TRUE(saw_reference_with_expected_literals); + EXPECT_TRUE(saw_retry_with_expected_literals); +} + TEST_F(HloIsolationTest, TestRunModuleUseFusionDebuggerOption) { DelegatingRunner test_runner(&this->test_runner()); @@ -902,6 +971,217 @@ ENTRY main { } } +TEST_F(HloIsolationTest, TestDumpHloOutputCallbacks) { + const char* hlo_text = R"( +HloModule TestModule + +ENTRY main { + a = f32[] parameter(0) + b = f32[] parameter(1) + add = f32[] add(a, b) + ROOT mul = f32[] multiply(add, b) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + + // 1. Normal execution with eval_literal_mutator + auto expected_literals = std::make_shared(); + auto mutator = [](absl::string_view hlo_name, Literal* literal) { + if (hlo_name == "add") { + *static_cast(literal->untyped_data()) += 10.0f; + } + }; + + std::vector callbacks = + CreateDumpHloOutputCallbacks(module.get(), expected_literals, mutator); + + for (const auto& cb : callbacks) { + Literal dummy = LiteralUtil::CreateR0(5.0f); + std::vector> lits = { + std::make_shared(std::move(dummy))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + + EXPECT_TRUE(expected_literals->contains("add")); + EXPECT_TRUE(expected_literals->contains("mul")); + EXPECT_FLOAT_EQ(expected_literals->at("add")->Get({}), 15.0f); + EXPECT_FLOAT_EQ(expected_literals->at("mul")->Get({}), 5.0f); + + // 2. Edge cases: empty/null literals and null expected_literals map + expected_literals->clear(); + std::vector> empty_lits; + callbacks[0].callback(/*replica_id=*/0, /*partition_id=*/0, empty_lits); + std::vector> null_lits = {nullptr}; + callbacks[0].callback(/*replica_id=*/0, /*partition_id=*/0, null_lits); + EXPECT_TRUE(expected_literals->empty()); + + std::vector null_map_callbacks = + CreateDumpHloOutputCallbacks(module.get(), /*expected_literals=*/nullptr, + nullptr); + Literal valid_lit = LiteralUtil::CreateR0(1.0f); + std::vector> valid_lits = { + std::make_shared(std::move(valid_lit))}; + null_map_callbacks[0].callback(/*replica_id=*/0, /*partition_id=*/0, + valid_lits); +} + +TEST_F(HloIsolationTest, TestComparisonHloOutputCallbacks) { + const char* hlo_text = R"( +HloModule TestModule + +ENTRY main { + a = f32[2,2] parameter(0) + ROOT add = f32[2,2] add(a, a) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ref_module, + ParseAndReturnVerifiedModule(hlo_text)); + + absl::flat_hash_map> ref_groups; + for (const auto* computation : ref_module->computations()) { + for (const auto* instruction : computation->MakeInstructionPostOrder()) { + GroupKey key(instruction->opcode(), instruction->shape().ToString()); + ref_groups[key].push_back(std::string(instruction->name())); + } + } + + ModuleIsolationOptions options; + options.abs_error_bound = 1e-4; + options.rel_error_bound = 1e-4; + auto result_mutex = std::make_shared(); + + // 1. Match & Eviction: matching literal removes expected literal and records + // no error + { + auto expected_literals = std::make_shared(); + (*expected_literals)["add"] = std::make_shared( + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}})); + HloIsolationTestResult test_result; + + std::unique_ptr clone = module->Clone(); + std::vector callbacks = + CreateComparisonHloOutputCallbacks(clone.get(), ref_groups, + expected_literals, *module, options, + result_mutex, &test_result); + + for (const auto& cb : callbacks) { + Literal actual = + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}}); + std::vector> lits = { + std::make_shared(std::move(actual))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + + EXPECT_FALSE(expected_literals->contains("add")); + EXPECT_EQ(test_result.numeric_checks_size(), 0); + } + + // 2. Mismatch: differing literal records numeric check failure and evicts + { + auto expected_literals = std::make_shared(); + (*expected_literals)["add"] = std::make_shared( + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}})); + HloIsolationTestResult test_result; + + std::unique_ptr clone = module->Clone(); + std::vector callbacks = + CreateComparisonHloOutputCallbacks(clone.get(), ref_groups, + expected_literals, *module, options, + result_mutex, &test_result); + + ::testing::TestPartResultArray failures; + { + ::testing::ScopedFakeTestPartResultReporter reporter( + ::testing::ScopedFakeTestPartResultReporter::INTERCEPT_ALL_THREADS, + &failures); + for (const auto& cb : callbacks) { + Literal mismatch = + LiteralUtil::CreateR2({{100.0f, 2.0f}, {3.0f, 4.0f}}); + std::vector> lits = { + std::make_shared(std::move(mismatch))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + } + + EXPECT_FALSE(expected_literals->contains("add")); + EXPECT_GT(failures.size(), 0); + ASSERT_EQ(test_result.numeric_checks_size(), 1); + EXPECT_EQ(test_result.numeric_checks(0).name(), "FusionDebugger:add"); + EXPECT_TRUE(test_result.numeric_checks(0).has_top_mismatch()); + } + + // 3. Incompatible Shape: logs warning and returns early without recording + // mismatch + { + auto expected_literals = std::make_shared(); + (*expected_literals)["add"] = std::make_shared( + LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f, 4.0f})); + HloIsolationTestResult test_result; + + std::unique_ptr clone = module->Clone(); + std::vector callbacks = + CreateComparisonHloOutputCallbacks(clone.get(), ref_groups, + expected_literals, *module, options, + result_mutex, &test_result); + + for (const auto& cb : callbacks) { + Literal actual = + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}}); + std::vector> lits = { + std::make_shared(std::move(actual))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + + EXPECT_EQ(test_result.numeric_checks_size(), 0); + } + + // 4. Edge cases: missing key, null map, empty/null literals + { + HloIsolationTestResult test_result; + auto empty_expected_literals = std::make_shared(); + std::unique_ptr clone = module->Clone(); + std::vector callbacks = + CreateComparisonHloOutputCallbacks(clone.get(), ref_groups, + empty_expected_literals, *module, + options, result_mutex, &test_result); + + for (const auto& cb : callbacks) { + Literal actual = + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}}); + std::vector> lits = { + std::make_shared(std::move(actual))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + EXPECT_EQ(test_result.numeric_checks_size(), 0); + + // Null expected_literals map + std::unique_ptr clone_null = module->Clone(); + std::vector callbacks_null = + CreateComparisonHloOutputCallbacks(clone_null.get(), ref_groups, + /*expected_literals=*/nullptr, + *module, options, result_mutex, + &test_result); + for (const auto& cb : callbacks_null) { + Literal actual = + LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}}); + std::vector> lits = { + std::make_shared(std::move(actual))}; + cb.callback(/*replica_id=*/0, /*partition_id=*/0, lits); + } + + // Null / empty literals in callback + std::vector> empty_lits; + callbacks_null[0].callback(/*replica_id=*/0, /*partition_id=*/0, + empty_lits); + std::vector> null_lits = {nullptr}; + callbacks_null[0].callback(/*replica_id=*/0, /*partition_id=*/0, null_lits); + EXPECT_EQ(test_result.numeric_checks_size(), 0); + } +} + TEST_F(HloIsolationTest, TestPopulateNumericCheckMismatches) { NumericCheck numeric_check; @@ -942,97 +1222,6 @@ TEST_F(HloIsolationTest, TestPopulateNumericCheckMismatches) { EXPECT_DOUBLE_EQ(numeric_check.top_mismatch().rel_error(), 2.5); } -TEST(FusionDebuggerTest, DirUsesUndeclaredOutputsDir) { - // Save environment variable - const char* original_env = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); - std::string original_val = original_env ? original_env : ""; - - // Set custom undeclared outputs dir - std::string custom_dir = "/some/custom/undeclared/outputs/dir"; - tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", custom_dir.c_str(), - /*overwrite=*/1); - - EXPECT_EQ(GetFusionDebuggerDir(), custom_dir); - - // Restore environment variable - if (!original_val.empty()) { - tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", original_val.c_str(), - /*overwrite=*/1); - } else { - tsl::unsetenv("TEST_UNDECLARED_OUTPUTS_DIR"); - } -} - -TEST(FusionDebuggerTest, FilePathUsesUndeclaredOutputsDir) { - // Save environment variable - const char* original_env = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); - std::string original_val = original_env ? original_env : ""; - - // Set custom undeclared outputs dir - std::string custom_dir = "/some/custom/undeclared/outputs/dir"; - tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", custom_dir.c_str(), - /*overwrite=*/1); - - EXPECT_EQ( - GetFusionDebuggerFilePath("my_op"), - tsl::io::JoinPath(custom_dir, "fusion-debugger-reference-my_op.bin")); - - // Restore environment variable - if (!original_val.empty()) { - tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", original_val.c_str(), - /*overwrite=*/1); - } else { - tsl::unsetenv("TEST_UNDECLARED_OUTPUTS_DIR"); - } -} - -TEST(FusionDebuggerTest, CleanUpAndGetLeftoverFiles) { - // We can write to the directory from GetFusionDebuggerDir(). - std::string debugger_dir = GetFusionDebuggerDir(); - - // Make sure it is cleaned up before starting - CleanUpAllFusionDebuggerFiles(); - EXPECT_THAT(GetLeftoverFusionDebuggerFiles(), IsEmpty()); - - // Create a debug file - std::string file_path = GetFusionDebuggerFilePath("test_cleanup_op"); - - // Write a dummy string to file - tsl::Env* env = tsl::Env::Default(); - ASSERT_OK(tsl::WriteStringToFile(env, file_path, "dummy data")); - - // Verify it exists in leftover files and via filesystem - EXPECT_THAT(GetLeftoverFusionDebuggerFiles(), Contains(file_path)); - - // Clean up - CleanUpAllFusionDebuggerFiles(); - - // Verify it no longer exists - EXPECT_THAT(GetLeftoverFusionDebuggerFiles(), IsEmpty()); - EXPECT_THAT(env->FileExists(file_path), - StatusIs(absl::StatusCode::kNotFound)); -} - -TEST(FusionDebuggerTest, DestructorCleansUpAllFiles) { - // Clear any existing leftover files first - CleanUpAllFusionDebuggerFiles(); - EXPECT_THAT(GetLeftoverFusionDebuggerFiles(), IsEmpty()); - - std::string file_path = GetFusionDebuggerFilePath("cleanup_destructor_test"); - tsl::Env* env = tsl::Env::Default(); - - { - absl::Cleanup cleanup = [] { CleanUpAllFusionDebuggerFiles(); }; - ASSERT_OK(tsl::WriteStringToFile(env, file_path, "test data")); - EXPECT_OK(env->FileExists(file_path)); - } - - // Destruction of cleanup should delete the file - EXPECT_THAT(env->FileExists(file_path), - StatusIs(absl::StatusCode::kNotFound)); - EXPECT_THAT(GetLeftoverFusionDebuggerFiles(), IsEmpty()); -} - TEST_F(HloIsolationTest, PopulateMismatchAnnotations_Basic) { const absl::string_view hlo_string = R"hlo( HloModule test_module diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator.cc b/third_party/xla/xla/tsl/framework/bfc_allocator.cc index d3e6dd7abd2b47..51de93e13aa239 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator.cc @@ -887,6 +887,8 @@ void BFCAllocator::FinishChunkAllocation(Chunk* chunk, size_t num_bytes) { } stats_.peak_bytes_in_use = std::max(stats_.peak_bytes_in_use, stats_.bytes_in_use); + stats_.peak_allocated_bytes = std::max( + stats_.peak_allocated_bytes, stats_.bytes_in_use + stats_.bytes_reserved); stats_.largest_alloc_size = std::max(stats_.largest_alloc_size, chunk->size); @@ -1609,6 +1611,7 @@ bool BFCAllocator::ClearStats() { absl::MutexLock l(mutex_); stats_.num_allocs = 0; stats_.peak_bytes_in_use = stats_.bytes_in_use; + stats_.peak_allocated_bytes = stats_.bytes_in_use + stats_.bytes_reserved; stats_.largest_alloc_size = 0; return true; } diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc index b823e5b4a89339..51284bb6e38055 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc @@ -724,6 +724,55 @@ TEST(BFCAllocatorTest, SpatialUnderContention) { EXPECT_EQ(failures.load(std::memory_order_relaxed), 0); } +TEST(BFCAllocatorTest, GetAndClearMemoryStats) { + BFCAllocator alloc(std::make_unique(), + /*total_memory=*/256 << 20, /*name=*/"test_stats", + BFCAllocator::Options{}); + + std::optional initial_stats = alloc.GetStats(); + ASSERT_TRUE(initial_stats.has_value()); + ASSERT_EQ(initial_stats->bytes_in_use, 0); + ASSERT_EQ(initial_stats->peak_bytes_in_use, 0); + ASSERT_EQ(initial_stats->peak_allocated_bytes, 0); + + const size_t kAllocSize1 = 1024; + void* ptr1 = alloc.AllocateRaw(kAlignment, kAllocSize1); + ASSERT_NE(ptr1, nullptr); + + std::optional stats_after_alloc1 = alloc.GetStats(); + ASSERT_TRUE(stats_after_alloc1.has_value()); + ASSERT_EQ(stats_after_alloc1->bytes_in_use, 1024); + ASSERT_EQ(stats_after_alloc1->peak_bytes_in_use, 1024); + ASSERT_EQ(stats_after_alloc1->peak_allocated_bytes, 1024); + + const size_t kAllocSize2 = 2048; + void* ptr2 = alloc.AllocateRaw(kAlignment, kAllocSize2); + ASSERT_NE(ptr2, nullptr); + + std::optional stats_after_alloc2 = alloc.GetStats(); + ASSERT_TRUE(stats_after_alloc2.has_value()); + ASSERT_EQ(stats_after_alloc2->bytes_in_use, 3072); + ASSERT_EQ(stats_after_alloc2->peak_bytes_in_use, 3072); + ASSERT_EQ(stats_after_alloc2->peak_allocated_bytes, 3072); + + alloc.DeallocateRaw(ptr2); + + std::optional stats_after_free = alloc.GetStats(); + ASSERT_TRUE(stats_after_free.has_value()); + ASSERT_EQ(stats_after_free->bytes_in_use, 1024); + ASSERT_EQ(stats_after_free->peak_bytes_in_use, 3072); + ASSERT_EQ(stats_after_free->peak_allocated_bytes, 3072); + + ASSERT_TRUE(alloc.ClearStats()); + std::optional stats_after_clear = alloc.GetStats(); + ASSERT_TRUE(stats_after_clear.has_value()); + EXPECT_EQ(stats_after_clear->bytes_in_use, 1024); + EXPECT_EQ(stats_after_clear->peak_bytes_in_use, 1024); + EXPECT_EQ(stats_after_clear->peak_allocated_bytes, 1024); + + alloc.DeallocateRaw(ptr1); +} + //===----------------------------------------------------------------------===// // Performance benchmarks. //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/tsl/framework/cpu_allocator_impl.cc b/third_party/xla/xla/tsl/framework/cpu_allocator_impl.cc index 7f0069362ba1d2..4b6f9c49eeb431 100644 --- a/third_party/xla/xla/tsl/framework/cpu_allocator_impl.cc +++ b/third_party/xla/xla/tsl/framework/cpu_allocator_impl.cc @@ -144,6 +144,7 @@ class CPUAllocator : public Allocator { absl::MutexLock l(mu_); stats_.num_allocs = 0; stats_.peak_bytes_in_use = stats_.bytes_in_use; + stats_.peak_allocated_bytes = stats_.bytes_in_use + stats_.bytes_reserved; stats_.largest_alloc_size = 0; return true; } @@ -183,6 +184,9 @@ class CPUAllocator : public Allocator { stats_.bytes_in_use += alloc_size; stats_.peak_bytes_in_use = std::max(stats_.peak_bytes_in_use, stats_.bytes_in_use); + stats_.peak_allocated_bytes = + std::max(stats_.peak_allocated_bytes, + stats_.bytes_in_use + stats_.bytes_reserved); stats_.largest_alloc_size = std::max(stats_.largest_alloc_size, alloc_size);