diff --git a/ci/official/envs/linux_x86_cuda b/ci/official/envs/linux_x86_cuda index 5e6ce4d7bef58d..dd93e2f4ded292 100644 --- a/ci/official/envs/linux_x86_cuda +++ b/ci/official/envs/linux_x86_cuda @@ -14,7 +14,7 @@ # ============================================================================== source ci/official/envs/linux_x86 export TF_FORCE_GPU_ALLOW_GROWTH=true -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true --local_test_jobs=16" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow" diff --git a/ci/official/envs/linux_x86_cuda13_nvcc b/ci/official/envs/linux_x86_cuda13_nvcc index bc8fc063157615..a8306754723f4f 100644 --- a/ci/official/envs/linux_x86_cuda13_nvcc +++ b/ci/official/envs/linux_x86_cuda13_nvcc @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== source ci/official/envs/linux_x86 -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version --local_test_jobs=16" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda_13_nvcc TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow_cuda13" diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index 435b63a90d65b5..a2acd13d151f4a 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -47,6 +47,7 @@ package_group( "//third_party/iree/...", "//third_party/odml/infra/...", "//third_party/odml/litert/...", + "//third_party/py/ai_edge_jax/...", "//waymo/accelerator/alpine/tools/...", "//waymo/ml/compiler/mlir/...", ], diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 5ed120ebbeade1..5b70d4b2bb3362 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -299,12 +299,26 @@ cc_library( "flatbuffer_to_mlir.h", ], deps = [ + "//tensorflow/compiler/mlir:op_or_arg_name_mapper", + "//tensorflow/compiler/mlir/lite:flatbuffer_export", "//tensorflow/compiler/mlir/lite:flatbuffer_import", + "//tensorflow/compiler/mlir/lite:tensorflow_lite", + "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps", + "//tensorflow/compiler/mlir/tensorflow", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:BytecodeWriter", + "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Parser", + "@llvm-project//mlir:QuantOps", "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", + "@stablehlo//:stablehlo_ops", + "@stablehlo//:vhlo_ops", ], ) diff --git a/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi b/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi index 401c020590f898..8125d0b9c17a4a 100644 --- a/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi +++ b/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi @@ -17,6 +17,7 @@ def Convert(model_flags_proto_txt_raw: object, converter_flags_proto_txt_raw: ob def ConvertMlirBytecode(converter_flags_proto_txt_raw: object, model_dir_txt_raw: object, output_file_path_raw: object) -> object: ... def ExperimentalMlirQuantizeModel(input_contents_txt_raw: object, disable_per_channel: bool = ..., fully_quantize: bool = ..., inference_type: int = ..., input_data_type: int = ..., output_data_type: int = ..., enable_numeric_verify: bool = ..., enable_whole_model_verify: bool = ..., op_blocklist: object = ..., node_blocklist: object = ..., enable_variable_quantization: bool = ..., disable_per_channel_for_dense_layers: bool = ..., debug_options_proto_txt_raw: object = ...) -> object: ... def ExperimentalMlirSparsifyModel(input_contents_txt_raw: object) -> object: ... -def FlatBufferToMlir(arg0: str, arg1: bool) -> str: ... +def FlatBufferToMlir(model: str, input_is_filepath: bool = ..., bytecode: bool = ..., cl_options: list[str] = ...) -> object: ... +def MlirToFlatBuffer(mlir: str, input_is_filepath: bool = ..., emit_builtin_tflite_ops: bool = ..., emit_select_tf_ops: bool = ..., emit_custom_ops: bool = ..., emit_stablehlo_ops: bool = ...) -> bytes: ... def RegisterCustomOpdefs(custom_opdefs_txt_raw: object) -> object: ... def RetrieveCollectedErrors() -> list: ... diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api.cc b/tensorflow/compiler/mlir/lite/python/converter_python_api.cc index 00f18729a4ec48..a653a6d30de321 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api.cc +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api.cc @@ -16,10 +16,10 @@ limitations under the License. #include -#include #include #include #include +#include #include #include "absl/container/flat_hash_set.h" @@ -517,8 +517,20 @@ std::vector RetrieveCollectedErrors() { } std::string FlatBufferFileToMlir(const std::string& model, - bool input_is_filepath) { - return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath); + bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath, bytecode, + cl_options); +} + +std::string MlirToFlatBufferFile(const std::string& mlir, + bool input_is_filepath, + bool emit_builtin_tflite_ops, + bool emit_select_tf_ops, bool emit_custom_ops, + bool emit_stablehlo_ops) { + return ::tensorflow::MlirToFlatBufferFile( + mlir, input_is_filepath, emit_builtin_tflite_ops, emit_select_tf_ops, + emit_custom_ops, emit_stablehlo_ops); } PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api.h b/tensorflow/compiler/mlir/lite/python/converter_python_api.h index 4ed970e9df0cb9..e53467d3e4586a 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api.h +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api.h @@ -64,8 +64,17 @@ PyObject* RegisterCustomOpdefs(PyObject* list); std::vector RetrieveCollectedErrors(); // Returns MLIR string dump of the given Flatbuffer model. -std::string FlatBufferFileToMlir(const std::string& model, - bool input_is_filepath); +std::string FlatBufferFileToMlir( + const std::string& model, bool input_is_filepath, bool bytecode = false, + const std::vector& cl_options = {}); + +// Converts MLIR (text or bytecode) to a TFLite Flatbuffer. +std::string MlirToFlatBufferFile(const std::string& mlir, + bool input_is_filepath, + bool emit_builtin_tflite_ops = true, + bool emit_select_tf_ops = false, + bool emit_custom_ops = true, + bool emit_stablehlo_ops = false); // Convert slim model to TfLite flatbuffer streamed directly to a file. PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc b/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc index 16af603063a0f4..57198b7ca87339 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include "pybind11/pybind11.h" // from @pybind11 +#include "pybind11/stl.h" // from @pybind11 #include "tensorflow/compiler/mlir/lite/python/converter_python_api.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h" #include "tensorflow/python/lib/core/pybind11_lib.h" @@ -119,12 +120,39 @@ PYBIND11_MODULE(_pywrap_converter_api, m, py::mod_gil_not_used()) { )pbdoc"); m.def( "FlatBufferToMlir", - [](const std::string& model, bool input_is_filepath) { - return tflite::FlatBufferFileToMlir(model, input_is_filepath); + [](const std::string& model, bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + std::string res = tflite::FlatBufferFileToMlir(model, input_is_filepath, + bytecode, cl_options); + if (bytecode) { + return py::object(py::bytes(res)); + } else { + return py::object(py::str(res)); + } }, + py::arg("model"), py::arg("input_is_filepath") = false, + py::arg("bytecode") = false, + py::arg("cl_options") = std::vector(), R"pbdoc( Returns MLIR dump of the given TFLite model. )pbdoc"); + m.def( + "MlirToFlatBuffer", + [](const std::string& mlir, bool input_is_filepath, + bool emit_builtin_tflite_ops, bool emit_select_tf_ops, + bool emit_custom_ops, bool emit_stablehlo_ops) { + std::string res = tflite::MlirToFlatBufferFile( + mlir, input_is_filepath, emit_builtin_tflite_ops, + emit_select_tf_ops, emit_custom_ops, emit_stablehlo_ops); + return py::bytes(res); + }, + py::arg("mlir"), py::arg("input_is_filepath") = false, + py::arg("emit_builtin_tflite_ops") = true, + py::arg("emit_select_tf_ops") = false, py::arg("emit_custom_ops") = true, + py::arg("emit_stablehlo_ops") = false, + R"pbdoc( + Converts MLIR (text or bytecode) into a TFLite FlatBuffer binary tensor. + )pbdoc"); m.def( "ConvertMlirBytecode", [](py::object converter_flags_proto_txt_raw, py::object model_dir_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc index b880df7f74a3ca..e80b33d18fba3f 100644 --- a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc +++ b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc @@ -18,21 +18,40 @@ limitations under the License. #include #include +#include "absl/base/attributes.h" +#include "absl/base/const_init.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "llvm/Support/CommandLine.h" #include "llvm/Support/LogicalResult.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/SMLoc.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/raw_ostream.h" +#include "mlir/Bytecode/BytecodeWriter.h" // from @llvm-project +#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/Quant.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/Diagnostics.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/IR/Verifier.h" // from @llvm-project +#include "mlir/Parser/Parser.h" // from @llvm-project #include "mlir/Support/FileUtilities.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo +#include "stablehlo/dialect/VhloOps.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h" #include "tensorflow/compiler/mlir/lite/flatbuffer_import.h" +#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" +#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h" +#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h" +#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h" namespace tensorflow { namespace { @@ -53,8 +72,33 @@ static mlir::OwningOpRef FlatBufferFileToMlirTranslation( } // namespace std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, - bool input_is_filepath) { - // referred logic from mlir::mlirTranslateMain(). + bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + ABSL_CONST_INIT static absl::Mutex cl_mutex(absl::kConstInit); + absl::MutexLock lock(&cl_mutex); + + // Reset options from any previous invocation. + llvm::cl::ResetAllOptionOccurrences(); + mlir::registerAsmPrinterCLOptions(); + + if (!cl_options.empty()) { + std::vector argv; + argv.reserve(cl_options.size() + 1); + argv.push_back("flatbuffer_to_mlir"); + for (const auto& opt : cl_options) { + argv.push_back(opt.c_str()); + } + std::string cl_errors; + llvm::raw_string_ostream cl_err_stream(cl_errors); + if (!llvm::cl::ParseCommandLineOptions( + argv.size(), argv.data(), "flatbuffer_to_mlir", &cl_err_stream)) { + cl_err_stream.flush(); + if (!cl_errors.empty()) { + llvm::errs() << "Failed to parse MLIR options: " << cl_errors << "\n"; + } + return ""; + } + } std::string errorMessage; std::unique_ptr input; @@ -78,17 +122,102 @@ std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, llvm::SourceMgr sourceMgr; sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc()); + std::string diagnostic_str; + llvm::raw_string_ostream diag_os(diagnostic_str); + mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os); + mlir::OwningOpRef module = FlatBufferFileToMlirTranslation(&sourceMgr, &context); - if (!module || failed(verify(*module))) return ""; + if (!module || failed(verify(*module))) { + diag_os.flush(); + if (!diagnostic_str.empty()) { + llvm::errs() << diagnostic_str << "\n"; + } + return ""; + } std::string mlir_output; llvm::raw_string_ostream output_stream(mlir_output); - // Dump MLIR with eliding large elements. - module->print( - output_stream, - mlir::OpPrintingFlags().useLocalScope().elideLargeElementsAttrs()); + if (bytecode) { + if (mlir::failed(mlir::writeBytecodeToFile(*module, output_stream))) { + llvm::errs() << "Failed to write MLIR bytecode.\n"; + return ""; + } + } else { + mlir::OpPrintingFlags flags; + module->print(output_stream, flags); + } + output_stream.flush(); return mlir_output; } +std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer, + bool input_is_filepath, + bool emit_builtin_tflite_ops, + bool emit_select_tf_ops, bool emit_custom_ops, + bool emit_stablehlo_ops) { + std::string errorMessage; + std::unique_ptr input; + if (input_is_filepath) { + input = mlir::openInputFile(mlir_file_or_buffer, &errorMessage); + if (!input) { + llvm::errs() << errorMessage << "\n"; + return ""; + } + } else { + input = + llvm::MemoryBuffer::getMemBuffer(mlir_file_or_buffer, "mlir", false); + if (!input) { + llvm::errs() << "Can't get llvm::MemoryBuffer\n"; + return ""; + } + } + + mlir::DialectRegistry registry; + registry.insert(); + mlir::RegisterAllTensorFlowDialects(registry); + + mlir::MLIRContext context(registry); + context.printOpOnDiagnostic(true); + + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc()); + + std::string diagnostic_str; + llvm::raw_string_ostream diag_os(diagnostic_str); + mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os); + + mlir::OwningOpRef module = + mlir::parseSourceFile(sourceMgr, &context); + if (!module || failed(verify(*module))) { + diag_os.flush(); + if (!diagnostic_str.empty()) { + llvm::errs() << diagnostic_str << "\n"; + } else { + llvm::errs() << "Failed to parse MLIR source.\n"; + } + return ""; + } + + std::string serialized_flatbuffer; + tensorflow::OpOrArgLocNameMapper op_or_arg_name_mapper; + tflite::FlatbufferExportOptions options; + options.converter_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops); + options.converter_flags.set_enable_select_tf_ops(emit_select_tf_ops); + options.converter_flags.set_allow_custom_ops(emit_custom_ops); + options.converter_flags.set_use_buffer_offset(true); + options.op_or_arg_name_mapper = &op_or_arg_name_mapper; + + if (!tflite::MlirToFlatBufferTranslateFunction( + *module, options, &serialized_flatbuffer, emit_stablehlo_ops)) { + llvm::errs() << "MlirToFlatBufferTranslateFunction failed.\n"; + return ""; + } + return serialized_flatbuffer; +} + } // namespace tensorflow diff --git a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h index 3164265f9764af..3a297d28c8328a 100644 --- a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h +++ b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h @@ -17,14 +17,24 @@ limitations under the License. #define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_ #include +#include namespace tensorflow { // Translates the given FlatBuffer filename or buffer into MLIR and returns -// translated MLIR as string. -std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, - bool input_is_filepath); - +// translated MLIR as string or bytecode. +std::string FlatBufferFileToMlir( + const std::string& model_file_or_buffer, bool input_is_filepath, + bool bytecode = false, const std::vector& cl_options = {}); + +// Translates the given MLIR filename or buffer into a TFLite FlatBuffer +// binary string. +std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer, + bool input_is_filepath, + bool emit_builtin_tflite_ops = true, + bool emit_select_tf_ops = false, + bool emit_custom_ops = true, + bool emit_stablehlo_ops = false); } // namespace tensorflow #endif // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_ diff --git a/tensorflow/compiler/mlir/lite/python/wrap_converter.py b/tensorflow/compiler/mlir/lite/python/wrap_converter.py index 47689da82d4d40..091be86c08cc73 100644 --- a/tensorflow/compiler/mlir/lite/python/wrap_converter.py +++ b/tensorflow/compiler/mlir/lite/python/wrap_converter.py @@ -98,6 +98,34 @@ def wrapped_retrieve_collected_errors(): return _pywrap_converter_api.RetrieveCollectedErrors() -def wrapped_flat_buffer_file_to_mlir(model, input_is_filepath): - """Wraps FlatBufferFileToMlir with lazy loader.""" - return _pywrap_converter_api.FlatBufferToMlir(model, input_is_filepath) +def wrapped_flat_buffer_file_to_mlir( + model, + input_is_filepath=False, + bytecode=False, + cl_options=None, +): + """Wraps FlatBufferToMlir with lazy loader.""" + if cl_options is None: + cl_options = [] + return _pywrap_converter_api.FlatBufferToMlir( + model, input_is_filepath, bytecode, cl_options + ) + + +def wrapped_mlir_to_flat_buffer( + mlir, + input_is_filepath=False, + emit_builtin_tflite_ops=True, + emit_select_tf_ops=False, + emit_custom_ops=True, + emit_stablehlo_ops=False, +): + """Wraps MlirToFlatBuffer with lazy loader.""" + return _pywrap_converter_api.MlirToFlatBuffer( + mlir, + input_is_filepath, + emit_builtin_tflite_ops, + emit_select_tf_ops, + emit_custom_ops, + emit_stablehlo_ops, + ) diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD index a9cf236970aac8..fc105c7c6fecec 100644 --- a/tensorflow/lite/python/BUILD +++ b/tensorflow/lite/python/BUILD @@ -284,7 +284,7 @@ py_test( ":util", #internal proto upb dep "//third_party/py/numpy", - "//tensorflow:tensorflow_py", + "//tensorflow:tensorflow_py_no_contrib", "//tensorflow/python/client:session", "//tensorflow/python/eager:context", "//tensorflow/python/eager:def_function", diff --git a/tensorflow/lite/python/analyzer.py b/tensorflow/lite/python/analyzer.py index 110d8014786b68..a147508ebe2e08 100644 --- a/tensorflow/lite/python/analyzer.py +++ b/tensorflow/lite/python/analyzer.py @@ -96,11 +96,22 @@ def analyze(model_path=None, input_is_filepath = False if kwargs.get("experimental_use_mlir", False): - print( - wrap_converter.wrapped_flat_buffer_file_to_mlir( - tflite_model, input_is_filepath - ) - ) + try: + mlir_text = wrap_converter.wrapped_flat_buffer_file_to_mlir( + tflite_model, + input_is_filepath, + cl_options=[ + "-mlir-print-local-scope", + "-mlir-elide-elementsattrs-if-larger=16", + ], + ) + except TypeError: + # Fallback for environments where wrapped_flat_buffer_file_to_mlir does + # not yet accept cl_options (e.g. OSS LiteRT using external TF). + mlir_text = wrap_converter.wrapped_flat_buffer_file_to_mlir( + tflite_model, input_is_filepath + ) + print(mlir_text) else: print( _analyzer_wrapper.ModelAnalyzer(tflite_model, input_is_filepath, diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 9d319e54b894b4..4bee63b60d07c0 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -198,6 +198,67 @@ def get_options(): return [str(option) for option in list(OpsSet)] +@convert_phase(Component.CONVERT_TF_TO_TFLITE_MODEL, SubComponent.UNSPECIFIED) +def flatbuffer_to_mlir( + model_content, + input_is_filepath=False, + bytecode=False, + cl_options=None, +): + """Converts a TFLite FlatBuffer model to MLIR string or bytecode. + + Args: + model_content: A TFLite FlatBuffer as bytes, or a path to a TFLite file if + input_is_filepath is True. + input_is_filepath: If True, model_content is treated as a file path. + bytecode: If True, returns MLIR bytecode (.mlirc / .mlirbc) as bytes. + Otherwise returns textual MLIR as a string. + cl_options: Sequence of MLIR command-line printing options to forward. + + Returns: + str if bytecode=False, bytes if bytecode=True. + """ + return wrap_converter.wrapped_flat_buffer_file_to_mlir( + model_content, + input_is_filepath=input_is_filepath, + bytecode=bytecode, + cl_options=cl_options, + ) + + +@convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.UNSPECIFIED) +def mlir_to_flatbuffer( + mlir_content, + input_is_filepath=False, + emit_builtin_tflite_ops=True, + emit_select_tf_ops=False, + emit_custom_ops=True, + emit_stablehlo_ops=False, +): + """Converts an MLIR source (textual or bytecode) to a TFLite FlatBuffer. + + Args: + mlir_content: MLIR source as a string (textual IR) or bytes (bytecode), or a + file path if input_is_filepath is True. + input_is_filepath: If True, mlir_content is treated as a file path. + emit_builtin_tflite_ops: Whether to emit builtin TFLite operations. + emit_select_tf_ops: Whether to emit Select TF operations (Flex ops). + emit_custom_ops: Whether to allow custom operations. + emit_stablehlo_ops: Whether to serialize StableHLO operations. + + Returns: + TFLite FlatBuffer as bytes. + """ + return wrap_converter.wrapped_mlir_to_flat_buffer( + mlir_content, + input_is_filepath=input_is_filepath, + emit_builtin_tflite_ops=emit_builtin_tflite_ops, + emit_select_tf_ops=emit_select_tf_ops, + emit_custom_ops=emit_custom_ops, + emit_stablehlo_ops=emit_stablehlo_ops, + ) + + @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.QUANTIZE) def mlir_quantize( input_data_str, diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 0c6dfa0c9adbb9..65a6416d1324de 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -46,8 +46,10 @@ from tensorflow.lite.python.convert import convert_saved_model as _convert_saved_model from tensorflow.lite.python.convert import ConverterError # pylint: disable=unused-import from tensorflow.lite.python.convert import deduplicate_readonly_buffers as _deduplicate_readonly_buffers +from tensorflow.lite.python.convert import flatbuffer_to_mlir as _flatbuffer_to_mlir from tensorflow.lite.python.convert import mlir_quantize as _mlir_quantize from tensorflow.lite.python.convert import mlir_sparsify as _mlir_sparsify +from tensorflow.lite.python.convert import mlir_to_flatbuffer as _mlir_to_flatbuffer from tensorflow.lite.python.convert import OpsSet from tensorflow.lite.python.convert import toco_convert # pylint: disable=unused-import from tensorflow.lite.python.convert_phase import Component @@ -3543,3 +3545,4 @@ def from_keras_model_file( return TFLiteConverter.from_keras_model_file( model_file, input_arrays, input_shapes, output_arrays ) + diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index da0c23f3ac5937..bc3e0440313956 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -2904,5 +2904,43 @@ def testDeprecatedOptionWarning(self, optimization): logging.root.removeHandler(handler) +class FlatbufferMlirTranslationTest(TestModels): + + def testFlatbufferToMlirAndBack(self): + saved_model_dir = os.path.join( + self.get_temp_dir(), 'simple_savedmodel_mlir' + ) + with ops.Graph().as_default(): + with session.Session() as sess: + in_tensor = array_ops.placeholder( + shape=[1, 16], dtype=dtypes.float32, name='input' + ) + out_tensor = in_tensor + 1.0 + inputs = {'x': in_tensor} + outputs = {'z': out_tensor} + saved_model.simple_save(sess, saved_model_dir, inputs, outputs) + converter = lite.TFLiteConverter.from_saved_model(saved_model_dir) + tflite_model = converter.convert() + self.assertIsNotNone(tflite_model) + + mlir_str = lite._flatbuffer_to_mlir(tflite_model) + self.assertIsInstance(mlir_str, str) + self.assertIn('module', mlir_str) + + mlir_elided = lite._flatbuffer_to_mlir( + tflite_model, cl_options=['-mlir-elide-elementsattrs-if-larger=8'] + ) + self.assertIsInstance(mlir_elided, str) + + mlir_bc = lite._flatbuffer_to_mlir(tflite_model, bytecode=True) + self.assertIsInstance(mlir_bc, bytes) + + reloaded_tflite = lite._mlir_to_flatbuffer(mlir_str) + self.assertIsNotNone(reloaded_tflite) + + reloaded_tflite_from_bc = lite._mlir_to_flatbuffer(mlir_bc) + self.assertIsNotNone(reloaded_tflite_from_bc) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh index 137897dd7c99c2..a00dcbc3f3404a 100755 --- a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -53,38 +53,31 @@ TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})" shift # ******************************************************************* -LOCK_DIR="${TF_LOCK_DIR:-/var/lock}" -mkdir -p "$LOCK_DIR" +mkdir -p /var/lock # Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU # slots to run a test at. # # Prefer to allocate 1 test per GPU over 4 tests on 1 GPU. # So, we iterate over TF_TESTS_PER_GPU first. -MAX_ATTEMPTS=30 -for attempt in $(seq 1 $MAX_ATTEMPTS); do - for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do - for i in `seq 0 $((TF_GPU_COUNT-1))`; do - exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exec {lock_fd}>&- - exit $return_code - fi - exec {lock_fd}>&- - done +for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exit $return_code + fi done - sleep 1 done -echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on, exiting with failure..." exit 1 diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index b1e2e8b7ff4ea0..fcff1f3810b3a6 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -15,4 +15,595 @@ # # This file is automatically generated by generate_patch tool. # Do not edit directly. +diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +--- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ++++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +@@ -20,6 +20,7 @@ + #include + + #include "llvm/ADT/DenseSet.h" ++#include "llvm/ADT/STLExtras.h" + #include "llvm/ADT/SmallVector.h" + #include "llvm/ADT/StringRef.h" + #include "mlir/IR/Attributes.h" +@@ -31,57 +32,53 @@ + namespace mlir { + namespace stablehlo { + +-static SmallVector> ++namespace { ++ ++struct ReindexedAxes { ++ SmallVector splitAxisSizes; ++ SmallVector groupedAxisIndices; ++}; ++ ++// Generates replica groups from the reshaped mesh axis sizes and the indices of ++// the communication axes using a Reshape-Transpose permutation. ++SmallVector> + flattenedReplicaGroupsFromTransposePermutation( +- const SmallVector& meshAxisNames, +- const SmallVector& commAxisNames, +- const llvm::DenseSet& commAxisSet, +- const SmallVector& axisSizes, +- const SmallVector& deviceIds, int64_t totalDevices) { +- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. ++ ArrayRef axisSizes, ArrayRef groupedAxisIndices, ++ ArrayRef deviceIds, int64_t totalDevices) { ++ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), ++ groupedAxisIndices.end()); + SmallVector transposeAxes; + // Non-grouped axes first +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (!commAxisSet.count(meshAxisNames[i])) { ++ for (size_t i = 0; i < axisSizes.size(); ++i) { ++ if (!groupedAxisSet.count(i)) { + transposeAxes.push_back(i); + } + } +- // Grouped axes +- for (const auto& name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- transposeAxes.push_back(i); +- break; +- } +- } +- } +- +- SmallVector transposedSizes(meshAxisNames.size()); +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { ++ // Grouped axes in the specified order ++ for (int64_t idx : groupedAxisIndices) { ++ transposeAxes.push_back(idx); ++ } ++ ++ SmallVector transposedSizes(axisSizes.size()); ++ for (size_t i = 0; i < axisSizes.size(); ++i) { + transposedSizes[i] = axisSizes[transposeAxes[i]]; + } + +- // Compute strides for original shape +- SmallVector originalStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ // Compute strides for reshaped shape ++ SmallVector originalStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; + } + + // Compute strides for transposed shape +- SmallVector transposedStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ SmallVector transposedStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; + } + +- // Generate chunks + int64_t numDevicesPerGroup = 1; +- for (auto name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- numDevicesPerGroup *= axisSizes[i]; +- break; +- } +- } ++ for (int64_t idx : groupedAxisIndices) { ++ numDevicesPerGroup *= axisSizes[idx]; + } + int64_t numGroups = totalDevices / numDevicesPerGroup; + +@@ -93,7 +90,7 @@ + for (int64_t j = 0; j < numDevicesPerGroup; ++j) { + int64_t linearTransposeIdx = i * numDevicesPerGroup + j; + int64_t originalIndex = 0; +- for (size_t k = 0; k < meshAxisNames.size(); ++k) { ++ for (size_t k = 0; k < axisSizes.size(); ++k) { + int64_t coord = + (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; + originalIndex += coord * originalStrides[transposeAxes[k]]; +@@ -102,9 +99,135 @@ + } + groups.push_back(std::move(group)); + } +- + return groups; + } ++ ++// Splits mesh axes based on sub-axis references and computes the corresponding ++// indices for the communication axes. ++FailureOr computeReindexedAxes(ArrayRef axesInMesh, ++ ArrayAttr commAxes, ++ Location loc) { ++ ReindexedAxes result; ++ ++ // Validate commAxes and verify that all mesh axes exist and have valid sizes. ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::dyn_cast(attr); ++ if (!shloAxisRef) { ++ return emitError(loc) << "expected AxisRefAttr in comm_axes"; ++ } ++ StringRef axisName = shloAxisRef.getName(); ++ bool found = false; ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ found = true; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ int64_t preSize = subAxisInfo.getPreSize(); ++ int64_t size = subAxisInfo.getSize(); ++ if (preSize < 1 || size < 1) { ++ return emitError(loc) ++ << "sub-axis pre_size and size must be at least 1"; ++ } ++ int64_t nextPreSize = preSize * size; ++ if (nextPreSize > meshAxis.getSize() || ++ meshAxis.getSize() % nextPreSize != 0) { ++ return emitError(loc) ++ << "sub-axis (pre_size * size) must divide mesh axis size"; ++ } ++ } ++ break; ++ } ++ } ++ if (!found) { ++ return emitError(loc) ++ << "axis '" << axisName << "' not found in mesh definition"; ++ } ++ } ++ ++ // Split each mesh axis according to the referenced subaxes. ++ struct SplitDim { ++ StringRef axisName; ++ int64_t preSize; ++ int64_t size; ++ int64_t dimIndex; ++ }; ++ SmallVector splitDims; ++ ++ for (auto meshAxis : axesInMesh) { ++ StringRef axisName = meshAxis.getName(); ++ int64_t axisSize = meshAxis.getSize(); ++ ++ if (axisSize == 1) { ++ int64_t dimIdx = result.splitAxisSizes.size(); ++ result.splitAxisSizes.push_back(1); ++ splitDims.push_back({axisName, /*preSize=*/1, /*size=*/1, dimIdx}); ++ continue; ++ } ++ ++ SmallVector preSizes = {1, axisSize}; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ if (shloAxisRef.getName() == axisName) { ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ preSizes.push_back(subAxisInfo.getPreSize()); ++ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); ++ } ++ } ++ } ++ ++ llvm::sort(preSizes); ++ preSizes.erase(llvm::unique(preSizes), preSizes.end()); ++ ++ for (size_t j = 0; j < preSizes.size() - 1; ++j) { ++ int64_t segPreSize = preSizes[j]; ++ int64_t segSize = preSizes[j + 1] / segPreSize; ++ int64_t dimIdx = result.splitAxisSizes.size(); ++ result.splitAxisSizes.push_back(segSize); ++ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); ++ } ++ } ++ ++ // Map each communication axis to its corresponding split dimension. ++ llvm::DenseSet groupedSet; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ StringRef axisName = shloAxisRef.getName(); ++ int64_t reqPreSize = 1; ++ int64_t reqSize = 0; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ reqPreSize = subAxisInfo.getPreSize(); ++ reqSize = subAxisInfo.getSize(); ++ } else { ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ reqSize = meshAxis.getSize(); ++ break; ++ } ++ } ++ } ++ ++ bool matched = false; ++ for (const auto& splitDim : splitDims) { ++ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && ++ splitDim.size == reqSize) { ++ if (!groupedSet.insert(splitDim.dimIndex).second) { ++ return emitError(loc) ++ << "Duplicate or overlapping communication axis: " << axisName; ++ } ++ result.groupedAxisIndices.push_back(splitDim.dimIndex); ++ matched = true; ++ break; ++ } ++ } ++ if (!matched) { ++ return emitError(loc) << "Invalid or overlapping communication axis on '" ++ << axisName << "'"; ++ } ++ } ++ ++ return result; ++} ++ ++} // namespace + + FailureOr>> flattenReplicaGroupMeshAxes( + Attribute meshAttr, ArrayAttr commAxes, Location loc) { +@@ -120,34 +243,13 @@ + if (!mesh) + return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); + +- auto axesInMesh = mesh.getAxes(); +- +- // Identify which axes are communication axes. +- llvm::SmallVector commAxisNames; +- llvm::DenseSet commAxisSet; +- for (auto attr : commAxes) { +- auto shloAxisRef = llvm::dyn_cast(attr); +- if (!shloAxisRef) { +- return emitError(loc) << "expected AxisRefAttr in comm_axes"; +- } +- if (shloAxisRef.getSubAxisInfo()) { +- return emitError(loc) << "Subaxes are not supported in " +- "flattenReplicaGroupMeshAxes"; +- } +- commAxisNames.push_back(shloAxisRef.getName()); +- commAxisSet.insert(shloAxisRef.getName()); +- } +- +- // Calculate total devices and axis sizes ++ FailureOr reindexedAxes = ++ computeReindexedAxes(mesh.getAxes(), commAxes, loc); ++ if (failed(reindexedAxes)) return failure(); + + int64_t totalDevices = 1; +- SmallVector axisSizes; +- SmallVector meshAxisNames; +- for (auto meshAxis : axesInMesh) { +- auto typedMeshAxis = llvm::cast(meshAxis); +- axisSizes.push_back(typedMeshAxis.getSize()); +- meshAxisNames.push_back(typedMeshAxis.getName()); +- totalDevices *= typedMeshAxis.getSize(); ++ for (auto meshAxis : mesh.getAxes()) { ++ totalDevices *= llvm::cast(meshAxis).getSize(); + } + + SmallVector deviceIds; +@@ -160,8 +262,8 @@ + } + + return flattenedReplicaGroupsFromTransposePermutation( +- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, +- totalDevices); ++ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, ++ deviceIds, totalDevices); + } + + } // namespace stablehlo +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir +--- stablehlo/stablehlo/tests/interpret/all_gather.mlir ++++ stablehlo/stablehlo/tests/interpret/all_gather.mlir +@@ -133,3 +133,41 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { ++ %result = "stablehlo.all_gather"(%operand) { ++ all_gather_dim = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<1xi64>) -> tensor<2xi64> ++ return %result : tensor<2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], ++ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) ++ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> ++ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir +--- stablehlo/stablehlo/tests/interpret/all_reduce.mlir ++++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir +@@ -135,3 +135,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.all_reduce"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, ++ channel_handle = #stablehlo.channel_handle ++ } : (tensor<1xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], ++ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir +--- stablehlo/stablehlo/tests/interpret/all_to_all.mlir ++++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir +@@ -172,3 +172,43 @@ + func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { ++ %result = "stablehlo.all_to_all"(%operand) { ++ split_dimension = 0 : i64, ++ concat_dimension = 1 : i64, ++ split_count = 2 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2x1xi64>) -> tensor<1x2xi64> ++ return %result : tensor<1x2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> ++ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> ++ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> ++ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> ++ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> ++ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> ++ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> ++ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], ++ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] ++ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, ++ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> ++ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, ++ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) ++ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +--- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ++++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +@@ -90,3 +90,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.reduce_scatter"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ scatter_dimension = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> ++ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> ++ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> ++ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> ++ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> ++ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> ++ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> ++ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], ++ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] ++ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +--- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir ++++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +@@ -5,7 +5,7 @@ + + // CHECK-LABEL: @all_reduce_rgv3 + func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -19,7 +19,7 @@ + + // CHECK-LABEL: @all_gather_rgv3 + func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_gather"(%arg0) { + all_gather_dim = 0 : i64, + replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> +@@ -30,7 +30,7 @@ + + // CHECK-LABEL: @all_to_all_rgv3 + func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_to_all"(%arg0) { + concat_dimension = 0 : i64, + split_dimension = 0 : i64, +@@ -44,7 +44,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh + func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -59,7 +59,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh_dev + func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -84,4 +84,56 @@ + } : (tensor<4xf32>) -> tensor<4xf32> + return %0 : tensor<4xf32> + } ++ ++ // CHECK-LABEL: @all_reduce_subaxis ++ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_1 ++ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_2 ++ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_size_one_axis ++ func.func @all_reduce_size_one_axis(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } + } diff --git a/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh index 137897dd7c99c2..a00dcbc3f3404a 100755 --- a/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -53,38 +53,31 @@ TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})" shift # ******************************************************************* -LOCK_DIR="${TF_LOCK_DIR:-/var/lock}" -mkdir -p "$LOCK_DIR" +mkdir -p /var/lock # Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU # slots to run a test at. # # Prefer to allocate 1 test per GPU over 4 tests on 1 GPU. # So, we iterate over TF_TESTS_PER_GPU first. -MAX_ATTEMPTS=30 -for attempt in $(seq 1 $MAX_ATTEMPTS); do - for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do - for i in `seq 0 $((TF_GPU_COUNT-1))`; do - exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exec {lock_fd}>&- - exit $return_code - fi - exec {lock_fd}>&- - done +for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exit $return_code + fi done - sleep 1 done -echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on, exiting with failure..." exit 1 diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc index 7c8f2862bea5f1..26997ce8df3b22 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc @@ -1960,12 +1960,10 @@ TEST_P(PrecisionTests, PrecisionCheck) { std::unique_ptr test_module, GetSimpleDotModule(kLhsOuterDim, kRhsOuterDim, kContractingDim, algorithm, backend)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector fake_arguments, - MakeFakeArguments(test_module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/23)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 23; + ASSERT_OK_AND_ASSIGN(std::vector fake_arguments, + MakeFakeArguments(test_module.get(), options)); // Ensure there are no negative arguments to avoid unbounded relative errors // due to subtracting two similarly large numbers. MakeNonNegative(fake_arguments); @@ -2030,13 +2028,10 @@ TEST_P(PrecisionTests, CheckPrecisionDegradationAlongKDimension) { TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr test_module, GetSimpleDotModule(kMSize, kNSize, k, algorithm, backend)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector fake_arguments, - MakeFakeArguments(test_module.get(), /*pseudo_random=*/ - true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/23)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 23; + ASSERT_OK_AND_ASSIGN(std::vector fake_arguments, + MakeFakeArguments(test_module.get(), options)); // Ensure there are no negative arguments to avoid unbounded relative errors // due to subtracting two similarly large numbers. MakeNonNegative(fake_arguments); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc index ba82ec58c82376..11c9f5ed824631 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc @@ -2571,8 +2571,10 @@ class TritonScaledDotTestBase : public TritonEmitterTest { absl::StatusOr> MakeScaledDotArguments( const HloModule* module) { std::minstd_rand0 engine; + FakeArgumentsOptions options; + options.engine = &engine; ABSL_ASSIGN_OR_RETURN(std::vector arguments, - MakeFakeArguments(module, &engine)); + MakeFakeArguments(module, options)); if (arguments.size() != 4) { return absl::InternalError(absl::StrCat( "Expected 4 scaled-dot arguments, got ", arguments.size())); diff --git a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc index cb9450a4362ce9..8cf1eaf1375ebe 100644 --- a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc @@ -44,12 +44,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR1({1, 2, 3}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -71,12 +69,10 @@ TEST_F(RaggedDotTest, NonContractingWithBatchDims) { lhs_ragged_dims={1}, rhs_group_dims={1} })"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=9. fake_arguments[2] = LiteralUtil::CreateR2({{4, 5}, {7, 2}, {6, 3}}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -98,12 +94,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR1({4, 2}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -125,12 +119,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR2({{1, 2, 3}, {3, 2, 1}}); EXPECT_TRUE(RunAndCompare(std::move(module), diff --git a/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc b/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc index af6e1098a2cb8f..1bbfaa4bdbddf3 100644 --- a/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include +#include #include #include "xla/error_spec.h" #include "xla/literal_util.h" @@ -45,11 +46,8 @@ ENTRY main { ROOT R = bf16[3072] reduce(prod, zero), dimensions={0}, to_apply=sum } )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(auto fake_arguments, MakeFakeArguments(module.get())); EXPECT_TRUE(RunAndCompare(std::move(module), LiteralUtil::MakePointers(fake_arguments), diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc index 7ea1ed00f02ba1..0c2fed894134ec 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc @@ -188,6 +188,8 @@ PjRtCApiClient::PjRtCApiClient( LOG(INFO) << "PjRtCApiClient created."; } +bool PjRtCApiClient::IsCApi() const { return true; } + void PjRtCApiClient::InitDevicesAndMemorySpaces() { // Initialize devices. PJRT_Client_Devices_Args devices_args; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h index 4b4da58f2ffdcc..c60c8df0da1613 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h @@ -372,6 +372,8 @@ class PjRtCApiClient : public PjRtClient { const PJRT_Api* c_api, PJRT_Client* c_client, std::unique_ptr<::pjrt::PJRT_KeyValueCallbackData> kv_callback_data); + bool IsCApi() const override; + int process_index() const override; int device_count() const override; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc index b7efe0487da9d0..a9c226fde55181 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc @@ -1137,5 +1137,12 @@ TEST(PjRtCApiClientTest, MakeCanonicalShapeForMemorySpace) { specific_layout.minor_to_major()); } +TEST(PjRtCApiClientTest, IsCApi) { + SetUpCpuPjRtApi(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr client, + GetCApiClient("cpu")); + EXPECT_TRUE(client->IsCApi()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc index ca8451d79ad792..95028d7f264cda 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc @@ -235,10 +235,10 @@ CreateSEGpuTopology(absl::string_view platform_name, } static se::StreamExecutor* GetFirstExecutor( - const std::vector>& devices) { - for (const auto& d : devices) { - if (auto* local_device_state = d.get()->local_device_state()) { - return local_device_state->executor(); + absl::Span> local_device_states) { + for (const auto& state : local_device_states) { + if (state != nullptr) { + return state->executor(); } } return nullptr; @@ -249,15 +249,16 @@ namespace { // Derives the platform version string (e.g. "cuda 12080") from the runtime // device description, independent of compile-time macros. static std::string GpuPlatformVersionFromDevices( + const StreamExecutorGpuRawClient* raw_client, absl::Span> devices) { for (const std::unique_ptr& device : devices) { - auto* se_device = device.get(); - LocalDeviceState* local_device_state = se_device->local_device_state(); - if (local_device_state == nullptr) { + auto local_device_state = + raw_client->GetLocalDeviceState(device->local_device_id()); + if (!local_device_state.ok()) { continue; } const se::DeviceDescription& desc = - local_device_state->executor()->GetDeviceDescription(); + local_device_state.value()->executor()->GetDeviceDescription(); const se::SemanticVersion v = desc.runtime_version(); const se::GpuComputeCapability& cc = desc.gpu_compute_capability(); if (cc.rocm_compute_capability() != nullptr) { @@ -306,13 +307,6 @@ void StreamExecutorGpuRawClient::UpdateCompileOptionsTopology( // helpers. namespace { -// Get the local device state for a given PjRtDevice. -absl::StatusOr GetLocalDeviceState(PjRtDevice* device) { - PjRtStreamExecutorDevice* pjrt_se_device = - absl::down_cast(device); - return pjrt_se_device->GetLocalDeviceState(); -} - // Creates a communicator for a cross-host transfer; used by the original // cross-host transfers API. absl::StatusOr> CreateTransferCommunicator( @@ -627,31 +621,29 @@ StreamExecutorGpuRawClient::CrossHostTransferBuffers( // Get the local_device_state and use it to schedule transfers. Fail // transfers early if we cannot get the local_device_state. - absl::StatusOr local_device_state = - absl::down_cast(device) - ->GetLocalDeviceState(); - if (!local_device_state.ok()) { - SetEventAsError(transfer_event, local_device_state.status()); + auto local_device_state_or = GetLocalDeviceState(device->local_device_id()); + if (!local_device_state_or.ok()) { + SetEventAsError(transfer_event, local_device_state_or.status()); continue; } + LocalDeviceState* local_device_state = *local_device_state_or; // Launch ScheduleTransfersOnLocalDevice on either the async dispatch thread // of the calling thread. - if ((*local_device_state)->async_dispatch_thread()) { - (*local_device_state) - ->async_dispatch_thread() - ->Schedule(tsl::WithCurrentContext( + if (local_device_state->async_dispatch_thread()) { + local_device_state->async_dispatch_thread()->Schedule( + tsl::WithCurrentContext( [this, local_device_state, device_id, transfer_dependencies, curr_transfer_specs = std::move(curr_transfer_specs), transfer_event = std::move(transfer_event)]() mutable { - ScheduleTransfersOnLocalDevice(*local_device_state, device_id, + ScheduleTransfersOnLocalDevice(local_device_state, device_id, std::move(transfer_event), std::move(transfer_dependencies), std::move(curr_transfer_specs)); })); } else { ScheduleTransfersOnLocalDevice( - *local_device_state, device_id, std::move(transfer_event), + local_device_state, device_id, std::move(transfer_event), transfer_dependencies, std::move(curr_transfer_specs)); } } @@ -1748,8 +1740,13 @@ absl::StatusOr BuildDistributedDevices( local_device_states_vec.push_back(std::move(it->second)); } + if (local_device != nullptr) { + CHECK_EQ(local_device->local_device_id().value(), + device_proto.local_device_ordinal()); + } auto device = std::make_unique( - device_proto.global_device_id(), local_device, device_proto.name(), + device_proto.global_device_id(), + /*is_addressable=*/local_device != nullptr, device_proto.name(), device_proto.vendor(), device_proto.compute_capability(), device_proto.core_count(), device_proto.device_memory_bytes_limit(), device_proto.shared_memory_per_block_optin(), @@ -1793,12 +1790,12 @@ absl::StatusOr BuildDistributedDevices( } StreamExecutorGpuDevice::StreamExecutorGpuDevice( - int id, LocalDeviceState* local_device_state, std::string device_kind, + int id, bool is_addressable, std::string device_kind, std::string device_vendor, std::string compute_capability, int core_count, int64_t device_memory_bytes_limit, int64_t shared_memory_per_block_optin, int local_device_id, int process_index, int process_index_in_partition, int partition_index, int numa_node, std::string fabric_uuid) - : PjRtStreamExecutorDevice(id, local_device_state, local_device_id, + : PjRtStreamExecutorDevice(id, is_addressable, local_device_id, process_index, process_index_in_partition, partition_index, std::move(device_kind)) { VLOG(1) << absl::StreamFormat( @@ -1894,7 +1891,8 @@ std::unique_ptr MakeStreamExecutorGpuClient( std::shared_ptr kv_store, std::shared_ptr topology, std::optional num_processes) { - std::string platform_version = GpuPlatformVersionFromDevices(devices); + std::string platform_version = + GpuPlatformVersionFromDevices(raw_client.get(), devices); PjRtPluginAttributes attrs; attrs.pjrt_c_api_major_version = 0; attrs.pjrt_c_api_minor_version = 0; @@ -2096,16 +2094,16 @@ absl::StatusOr> GetStreamExecutorGpuClient( ABSL_ASSIGN_OR_RETURN(std::shared_ptr gpu_topology, GpuTopology::FromProto(devices_and_topology.topology)); - auto se_gpu_topology = - CreateSEGpuTopology(pjrt_platform_name, std::move(gpu_topology), - GetFirstExecutor(devices_and_topology.devices)); + se::StreamExecutor* first_executor = + GetFirstExecutor(devices_and_topology.local_device_states); + auto se_gpu_topology = CreateSEGpuTopology( + pjrt_platform_name, std::move(gpu_topology), first_executor); auto raw_client = std::make_unique( tsl::Fingerprint64(pjrt_platform_name), std::move(devices_and_topology.local_device_states), std::move(allocator), xla_client, std::move(host_memory_allocator), options.should_stage_host_to_device_transfers, - /*async_work_runner=*/nullptr, - GetFirstExecutor(devices_and_topology.devices), kv_store, + /*async_work_runner=*/nullptr, first_executor, kv_store, preallocate_device_memory, options.abort_collectives_on_failure, std::move(gpu_run_options), std::move(memory_registration)); VLOG(1) << absl::StreamFormat( @@ -2156,16 +2154,16 @@ absl::StatusOr> GetSharedStreamExecutorGpuClient( ABSL_ASSIGN_OR_RETURN(auto gpu_topology, absl::StatusOr>( GpuTopology::FromProto(devices_and_topology.topology))); - auto se_gpu_topology = - CreateSEGpuTopology(platform_name, std::move(gpu_topology), - GetFirstExecutor(devices_and_topology.devices)); + se::StreamExecutor* first_executor = + GetFirstExecutor(devices_and_topology.local_device_states); + auto se_gpu_topology = CreateSEGpuTopology( + platform_name, std::move(gpu_topology), first_executor); auto raw_client = std::make_unique( tsl::Fingerprint64(platform_name), std::move(devices_and_topology.local_device_states), std::move(allocator), local_client, std::move(host_memory_allocator), /*should_stage_host_to_device_transfers=*/true, - /*async_work_runner=*/nullptr, - GetFirstExecutor(devices_and_topology.devices), kv_store, + /*async_work_runner=*/nullptr, first_executor, kv_store, /*cache_fabric_handles=*/false, /*abort_collectives_on_failure=*/false, std::move(gpu_run_options)); VLOG(1) << absl::StreamFormat( diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h index e79b15e8b6ed91..3467c593887e28 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h @@ -76,8 +76,8 @@ namespace xla { class StreamExecutorGpuDevice : public PjRtStreamExecutorDevice { public: - StreamExecutorGpuDevice(int id, LocalDeviceState* local_device_state, - std::string device_kind, std::string device_vendor, + StreamExecutorGpuDevice(int id, bool is_addressable, std::string device_kind, + std::string device_vendor, std::string compute_capability, int core_count, int64_t device_memory_bytes_limit, int64_t shared_memory_per_block_optin, diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc index b3c4969501fa5f..e31a33843a4801 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc @@ -1678,11 +1678,11 @@ absl::Status InterProcessCollectiveInitTestBody(int rank_id) { // executor's collective memory allocator into the selected collectives // backend (e.g. MORI ShmemMalloc). With inert backend stubs the allocation // may return null; we only log the outcome and do not fail the test. - auto* se_device = absl::down_cast( - client->addressable_devices()[0]); - TF_RET_CHECK(se_device != nullptr); - LocalDeviceState* local_device_state = se_device->local_device_state(); - TF_RET_CHECK(local_device_state != nullptr); + auto* se_client = absl::down_cast(client.get()); + TF_RET_CHECK(se_client != nullptr); + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device_state, + se_client->raw_client()->GetLocalDeviceState( + client->addressable_devices()[0]->local_device_id())); se::StreamExecutor* executor = local_device_state->executor(); constexpr uint64_t kCollectiveBytes = 1024; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index 593f3cb71a3733..86e4bd4d8095e9 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -229,13 +229,13 @@ TEST(StreamExecutorGpuClientTest, AsyncResultDefinitionEventUsesAsyncStream) { ASSERT_OK_AND_ASSIGN(auto definition, GetDefinitionStreamInfo(results[0][0].get())); - auto* se_device = absl::down_cast( - client->addressable_devices().front()); - intptr_t compute_stream = - reinterpret_cast(se_device->local_device_state() - ->compute_stream() - ->platform_specific_handle() - .stream); + auto* se_client = absl::down_cast(client.get()); + TF_ASSERT_OK_AND_ASSIGN( + LocalDeviceState * local_device_state, + se_client->raw_client()->GetLocalDeviceState( + client->addressable_devices().front()->local_device_id())); + intptr_t compute_stream = reinterpret_cast( + local_device_state->compute_stream()->platform_specific_handle().stream); EXPECT_NE(definition.stream, compute_stream); ASSERT_OK_AND_ASSIGN(auto literal, results[0][0]->ToLiteral().Await()); @@ -2568,12 +2568,11 @@ TEST(StreamExecutorGpuClientTest, EventCaching) { absl::down_cast(client.get()) ->async_work_runner(); const auto& device = client->addressable_devices()[0]; - // TODO(b/b/482307468) Switch to absl::down_cast after upgrade. - [[deprecated( - "remove after absl upgrade")]] LocalDeviceState* local_device_state = - absl::down_cast(device) - ->local_device_state(); - ASSERT_TRUE(local_device_state != nullptr); + TF_ASSERT_OK_AND_ASSIGN( + LocalDeviceState * local_device_state, + absl::down_cast(client.get()) + ->raw_client() + ->GetLocalDeviceState(device->local_device_id())); size_t sync_point0 = local_device_state->GetNextComputeStreamSyncPoint(); TF_ASSERT_OK_AND_ASSIGN(auto event0, local_device_state->GetEventForComputeStreamSyncPoint( diff --git a/third_party/xla/xla/pjrt/pjrt_client.cc b/third_party/xla/xla/pjrt/pjrt_client.cc index efca03ade466aa..28c34b4ca7ef62 100644 --- a/third_party/xla/xla/pjrt/pjrt_client.cc +++ b/third_party/xla/xla/pjrt/pjrt_client.cc @@ -119,6 +119,8 @@ PjRtMemorySpaceCApiDelegator::PjRtMemorySpaceCApiDelegator( PjRtBuffer::ExternalReference::~ExternalReference() = default; +bool PjRtClient::IsCApi() const { return false; } + absl::StatusOr PjRtClient::UnsafeBufferPointer( PjRtBuffer* buffer) { if (buffer->on_device_shape().IsTuple()) { diff --git a/third_party/xla/xla/pjrt/pjrt_client.h b/third_party/xla/xla/pjrt/pjrt_client.h index 27ed24b15259e3..8ad889abe2dabc 100644 --- a/third_party/xla/xla/pjrt/pjrt_client.h +++ b/third_party/xla/xla/pjrt/pjrt_client.h @@ -558,6 +558,9 @@ class PjRtClient { virtual ~PjRtClient() = default; + // Whether this client uses the C API. + virtual bool IsCApi() const; + // Return the process index of this client. Always 0 in single-process // settings. virtual int process_index() const = 0; diff --git a/third_party/xla/xla/pjrt/pjrt_client_test.cc b/third_party/xla/xla/pjrt/pjrt_client_test.cc index 7f6b62f30afb13..7343158f36e4d8 100644 --- a/third_party/xla/xla/pjrt/pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/pjrt_client_test.cc @@ -799,6 +799,11 @@ ENTRY RuntimeDonationDenialMustAliasFails() -> f32[2, 2] { TEST(PjRtClientTest, GetDefaultLayout) {} +TEST(PjRtClientTest, IsCApiTest) { + TF_ASSERT_OK_AND_ASSIGN(auto client, GetClient()); + EXPECT_FALSE(client->IsCApi()); +} + TEST(PjRtClientTest, ClearPeakMemory) { TF_ASSERT_OK_AND_ASSIGN(auto client, GetClient()); PjRtDevice* device = client->addressable_devices()[0]; diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index e680c677ef0429..f433928799737b 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -235,14 +235,6 @@ absl::string_view PjRtStreamExecutorDevice::platform_name() const { return client_->platform_name(); } -absl::StatusOr -PjRtStreamExecutorDevice::GetLocalDeviceState() const { - if (local_device_state_ != nullptr) { - return local_device_state_; - } - return InvalidArgument("Device %s is not a local device.", DebugString()); -} - absl::StatusOr DevicesToDeviceAssignment( absl::Span> devices) { if (devices.empty()) { @@ -427,10 +419,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( bool retry_on_oom, tsl::AsyncValueRef allocate_after) { CHECK(allocate_after == nullptr) << "allocate_after is not supported for PjRtStreamExecutorClient."; - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); PjRtMemorySpace* default_memory_space = device->default_memory_space().value_or(nullptr); auto layout_memory_space = Layout::kDefaultMemorySpace; @@ -477,10 +468,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBufferForExecute( PjRtMemorySpace* memory_space, size_t on_device_bytes_count, bool retry_on_oom) { - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto mem = RawSEDeviceMemory::CreateDelayedMemory(); return tsl::MakeRef( this, memory_space, local_device, std::move(mem), on_device_bytes_count); @@ -491,10 +481,9 @@ absl::StatusOr( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto raw_buffer = tsl::MakeRef( this, memory_space, local_device, tsl::AsyncValueRef(buffer_promise), @@ -614,10 +603,9 @@ PjRtStreamExecutorClient::AllocateLinearizeDest( absl::StatusOr> PjRtStreamExecutorRawClient::CreateLinkedEventPromise( PjRtMemorySpace* memory_space, absl::string_view debug_info) { - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto result = tsl::MakeRef( this, local_device, async_work_runner()); PjRtDeviceEventRef event = result->event().CopyRef(); @@ -637,10 +625,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::CreateDeviceEvent(PjRtMemorySpace* memory_space, Future<> dependency) { auto definition_event = BufferSequencingEvent::Create(async_work_runner()); - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); dependency.OnReady([definition_event = definition_event.CopyRef(), local_device, this](absl::Status status) mutable { if (!status.ok()) { @@ -704,8 +691,7 @@ PjRtStreamExecutorRawClient::ImportForeignMemory( CHECK_EQ(memory_space->devices().size(), 1); auto* device = memory_space->devices().front(); ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - tensorflow::down_cast(device) - ->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto buffer = RawSEDeviceMemory::CreateForeign( se::DeviceAddressBase(device_ptr, size), std::move(on_delete_callback)); @@ -720,8 +706,7 @@ PjRtStreamExecutorRawClient::CreateDeviceEventForStream( CHECK_EQ(memory_space->devices().size(), 1); auto* device = memory_space->devices().front(); ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - absl::down_cast(device) - ->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto definition_event = BufferSequencingEvent::Create(this->async_work_runner()); @@ -741,10 +726,7 @@ PjRtStreamExecutorExecutableLoadState::LoadRawExecutable( const ExecuteOptions& options, size_t host_callback_idx, xla::RunId run_id, DeviceAndAssignment device_and_assign, int attempt) { PjRtDevice* device = device_and_assign.device; - int device_ordinal = absl::down_cast(device) - ->local_device_state() - ->local_device_id() - .value(); + int device_ordinal = device->local_device_id().value(); auto se_executable = std::move(executable).Cast(); const CompileOptions& compile_options = se_executable->compile_options(); ABSL_ASSIGN_OR_RETURN(auto local_exec, @@ -762,30 +744,40 @@ PjRtStreamExecutorExecutableLoadState::LoadRawExecutable( std::move(se_executable), raw_client()); } +absl::Status PjRtStreamExecutorRawClient::TransferToInfeed( + LocalDeviceId local_device_id, const LiteralSlice& literal) { + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + GetLocalDeviceState(local_device_id)); + return NeverRunOnFiber(async_work_runner(), [&]() { + return local_device->client()->TransferToInfeedLocal( + literal, local_device->local_hardware_id().value()); + }); +} + +absl::Status PjRtStreamExecutorRawClient::TransferFromOutfeed( + LocalDeviceId local_device_id, MutableBorrowingLiteral literal) { + VLOG(1) << "PjRtStreamExecutorRawClient::TransferFromOutfeed"; + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + GetLocalDeviceState(local_device_id)); + return NeverRunOnFiber(async_work_runner(), [&]() { + return local_device->client()->TransferFromOutfeedLocal( + local_device->local_hardware_id().value(), literal); + }); +} + // Transfer the given literal to the infeed queue of the given local device. absl::Status PjRtStreamExecutorDevice::TransferToInfeed( const LiteralSlice& literal) { - // Only support infeed to local device. - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); - return NeverRunOnFiber( - tensorflow::down_cast(client_) - ->async_work_runner(), - [&]() { - return local_device->client()->TransferToInfeedLocal( - literal, local_device->local_hardware_id().value()); - }); + return absl::down_cast(client_) + ->raw_client() + ->TransferToInfeed(local_device_id(), literal); } absl::Status PjRtStreamExecutorDevice::TransferFromOutfeed( MutableBorrowingLiteral literal) { - VLOG(1) << "PjRtStreamExecutorDevice::TransferFromOutfeed"; - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); - return NeverRunOnFiber( - absl::down_cast(client_)->async_work_runner(), - [&]() { - return local_device->client()->TransferFromOutfeedLocal( - local_device->local_hardware_id().value(), literal); - }); + return absl::down_cast(client_) + ->raw_client() + ->TransferFromOutfeed(local_device_id(), literal); } void PjRtStreamExecutorDevice::AttachMemorySpace(PjRtMemorySpace* memory_space, @@ -846,7 +838,10 @@ PjRtStreamExecutorDevice::memory_space_by_kind_id(int id) const { absl::StatusOr PjRtStreamExecutorDevice::GetStreamForExternalReadyEvents() const { - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + absl::down_cast(client_) + ->raw_client() + ->GetLocalDeviceState(local_device_id())); se::Stream* stream = local_device->GetExternalReadyEventStream(); void* raw_stream = stream->platform_specific_handle().stream; if (raw_stream == nullptr) { @@ -1407,13 +1402,9 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( PjRtDeviceEventRefVector extra_deps, PjRtDeviceEventRefVector control_deps, bool is_predetermined_error, bool fill_future) && { const uint64_t start_time_usecs = tsl::Env::Default()->NowMicros(); - int device_ordinal = tensorflow::down_cast(device_) - ->local_device_state() - ->local_device_id() - .value(); + int device_ordinal = device_->local_device_id().value(); LocalDeviceState* device_state = - tensorflow::down_cast(device_) - ->local_device_state(); + raw_client_->device_state(device_->local_device_id()); const CompileOptions& compile_options = se_executable_->compile_options(); tsl::profiler::TraceMe trace([&] { @@ -1693,9 +1684,6 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( }; auto definition_event = [&]() -> PjRtDeviceEventRef { - LocalDeviceState* device_state = - tensorflow::down_cast(device) - ->local_device_state(); se::Stream* stream = device_state->compute_stream(); if (!result_buffer_or_status.ok()) { @@ -1805,8 +1793,7 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( void PjRtStreamExecutorClient::LaunchOnDevice( PjRtDevice* device, absl::AnyInvocable execute_fn) const { const LocalDeviceState& device_state = - *tensorflow::down_cast(device) - ->local_device_state(); + this->device_state(device->local_device_id().value()); device_state.execute_thread()->Schedule( tsl::WithCurrentContext(std::move(execute_fn))); } @@ -2190,8 +2177,7 @@ PjRtStreamExecutorClient::Load(std::shared_ptr executable, LoadInternal(std::move(executable), load_options, /*dump=*/false); for (const PjRtDevice* device : addressable_devices()) { LocalDeviceState* local_device_state = - tensorflow::down_cast(device) - ->local_device_state(); + raw_client()->device_state(device->local_device_id()); raw_client()->RecordMemoryStats(local_device_state); } return loaded_executable; @@ -2202,13 +2188,12 @@ bool PjRtStreamExecutorClient::IsHostMemoryPinned(const void* ptr, if (addressable_devices().empty()) { return false; } - auto* device = tensorflow::down_cast( - addressable_devices()[0]); - auto status_or_device_state = device->GetLocalDeviceState(); - if (!status_or_device_state.ok()) { + auto local_device_state = raw_client()->GetLocalDeviceState( + addressable_devices()[0]->local_device_id()); + if (!local_device_state.ok()) { return false; } - return status_or_device_state.value() + return local_device_state.value() ->compute_stream() ->parent() ->IsHostMemoryPinned(ptr, size); diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h index 359820015aedd9..81a3823a8f0adf 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h @@ -106,22 +106,20 @@ class StreamExecutorExecutable; class PjRtStreamExecutorDevice : public PjRtDevice { public: - PjRtStreamExecutorDevice(int id, LocalDeviceState* local_device_state, - int local_device_id, int process_index, - int process_index_in_partition, int partition_index, - std::string device_kind) - : local_device_id_(local_device_id), - local_hardware_id_(local_device_state - ? local_device_state->local_hardware_id() - : LocalChipId(-1)), - local_device_state_(local_device_state), + PjRtStreamExecutorDevice(int id, bool is_addressable, int local_device_id, + int process_index, int process_index_in_partition, + int partition_index, std::string device_kind, + LocalChipId local_hardware_id = LocalChipId(-1)) + : is_addressable_(is_addressable), + local_device_id_(local_device_id), + local_hardware_id_(local_hardware_id.value() != -1 + ? local_hardware_id + : (is_addressable ? LocalChipId(local_device_id) + : LocalChipId(-1))), description_(id, local_device_id_.value(), process_index, process_index_in_partition, partition_index, - std::move(device_kind)) { - if (local_device_state_ != nullptr) { - CHECK_EQ(local_device_state_->local_device_id(), local_device_id_); - } - } + std::move(device_kind)) {} + ~PjRtStreamExecutorDevice() override = default; // Must set client exactly once. @@ -155,7 +153,7 @@ class PjRtStreamExecutorDevice : public PjRtDevice { PjRtClient* client() const override { return client_; } - bool IsAddressable() const override { return local_device_state_ != nullptr; } + bool IsAddressable() const override { return is_addressable_; } LocalDeviceId local_device_id() const override { return local_device_id_; } @@ -166,16 +164,6 @@ class PjRtStreamExecutorDevice : public PjRtDevice { return attributes_; } - // If this is a device local to this host, returns a LocalDeviceState object - // that can be used to manipulate the device. Returns nullptr if the device is - // not local to this host. - LocalDeviceState* local_device_state() const { return local_device_state_; } - - // If this is a device local to this host, returns a LocalDeviceState object - // that can be used to manipulate the device. Returns an error if the device - // is not local to this host. - absl::StatusOr GetLocalDeviceState() const; - absl::Status TransferToInfeed(const LiteralSlice& literal) override; absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) override; @@ -201,9 +189,9 @@ class PjRtStreamExecutorDevice : public PjRtDevice { } private: + const bool is_addressable_; const LocalDeviceId local_device_id_; const LocalChipId local_hardware_id_; - LocalDeviceState* local_device_state_ = nullptr; PjRtStreamExecutorDeviceDescription description_; absl::flat_hash_map attributes_; PjRtClient* client_ = nullptr; @@ -263,6 +251,16 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { return it != local_device_states_by_id_.end() ? it->second : nullptr; } + absl::StatusOr GetLocalDeviceState( + LocalDeviceId local_device_id) const { + LocalDeviceState* state = device_state(local_device_id); + if (state == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "Device ", local_device_id.value(), " is not a local device.")); + } + return state; + } + gpu::GpuExecutableRunOptions* gpu_run_options() const { return gpu_run_options_.get(); } @@ -362,6 +360,12 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { absl::StatusOr CreateDeviceEventForStream( PjRtMemorySpace* memory_space, std::intptr_t stream) override; + absl::Status TransferToInfeed(LocalDeviceId local_device_id, + const LiteralSlice& literal) override; + + absl::Status TransferFromOutfeed(LocalDeviceId local_device_id, + MutableBorrowingLiteral literal) override; + virtual void UpdateCompileOptionsTopology( const PjRtTopologyDescription& topology, CompileOptions* options) const {} diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index 9d08e8d509e83c..baf79d2c5a3fe5 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -169,10 +169,9 @@ MakeTestPjRtStreamExecutorClient( std::unique_ptr gpu_run_options = nullptr, std::shared_ptr kv_store = nullptr) { se::StreamExecutor* first_executor = nullptr; - for (const auto& dev : devices) { - if (dev->IsAddressable() && dev->local_device_state() != nullptr && - dev->local_device_state()->compute_stream() != nullptr) { - first_executor = dev->local_device_state()->compute_stream()->parent(); + for (const auto& state : local_device_states) { + if (state != nullptr && state->compute_stream() != nullptr) { + first_executor = state->compute_stream()->parent(); break; } } @@ -209,7 +208,7 @@ absl::StatusOr> GetClient() { int local_device_id = local_device_states.back()->local_device_id().value(); std::vector> devices; devices.emplace_back(std::make_unique( - 0, local_device_states.back().get(), local_device_id, /*process_index=*/0, + 0, /*is_addressable=*/true, local_device_id, /*process_index=*/0, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); std::vector> memory_spaces; memory_spaces.emplace_back(std::make_unique( @@ -261,7 +260,7 @@ absl::StatusOr> GetClientWithDevices( /*allow_event_reuse=*/false, /*use_callback_stream=*/false)); int local_device_id = local_device_states.back()->local_device_id().value(); devices.emplace_back(std::make_unique( - i, local_device_states.back().get(), local_device_id, + i, /*is_addressable=*/true, local_device_id, /*process_index=*/0, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); memory_spaces.emplace_back(std::make_unique( @@ -271,7 +270,7 @@ absl::StatusOr> GetClientWithDevices( } for (int i = num_addressable_devices; i < total_devices; ++i) { devices.emplace_back(std::make_unique( - i, /*local_device_state=*/nullptr, /*local_device_id=*/-1, + i, /*is_addressable=*/false, /*local_device_id=*/-1, /*process_index=*/1, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); } @@ -485,7 +484,7 @@ TEST(PjRtStreamExecutorClientTest, ExecutePortableRemoteDevice) { *client, shape, [](XlaBuilder& builder) {}, compile_options)); auto remote_device = std::make_unique( - 1, /*local_device_state=*/nullptr, /*local_device_id=*/-1, + 1, /*is_addressable=*/false, /*local_device_id=*/-1, /*process_index=*/1, /*process_index_in_partition=*/1, /*partition_index=*/0, "cpu"); remote_device->SetClient(client.get()); diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc index 41afad3c6099ed..4cf23bec75c33f 100644 --- a/third_party/xla/xla/pjrt/transpose.cc +++ b/third_party/xla/xla/pjrt/transpose.cc @@ -1137,7 +1137,16 @@ absl::Status TransposePlan::Initialize() { ++pos_stride1a_in_a; } + is_identity_permutation_ = true; + for (int i = 0; i < permutation_.size(); ++i) { + if (permutation_[i] != i) { + is_identity_permutation_ = false; + break; + } + } + b_dims_ = Permute(a_dims_, permutation_); + ComputeStrides(elem_size_in_bytes_, b_dims_, b_tiling_, ldb_, ldb_tile_); // Find the innermost dimension of B that is stride 1 element. We know such a @@ -1377,7 +1386,17 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { loop_order.clear(); loop_order.reserve(remaining.size()); - auto soft_cost = [&](const Loop& l) { + // Computes a sort key for each candidate loop, returning a 3-tuple: + // + // Lexicographical ordering selects the smallest tuple first, placing that + // loop into an outer position (from slowest-varying to fastest-varying): + // 1. Contiguity category (non-contiguous loops placed in outer loops). + // 2. Primary cost: cache-line misses (higher miss cost placed in outer + // loops). + // 3. Secondary tie-breaker: when cache misses tie (e.g. both loops have + // strides >= 128B), the loop with larger minimum stride goes outer to + // favor better locality in at least one buffer for the inner loop. + auto soft_cost = [&](const Loop& l) -> std::tuple { int64_t a_stride = std::abs(l.lda); if (!inner_kernel_is_memcpy_ && l.is_inner_dim_in_a) { a_stride *= inner_block_elems_ * outer_block_elems_a_; @@ -1387,21 +1406,41 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { b_stride *= inner_block_elems_ * outer_block_elems_b_; } - double stride; switch (chunk_contiguity_) { case ChunkContiguity::kOutput: - stride = b_stride; - return std::make_tuple(0, -stride); + return std::make_tuple(0, -b_stride, 0.0); case ChunkContiguity::kInput: - stride = a_stride; - return std::make_tuple(0, -stride); - case ChunkContiguity::kNone: - // Add a small penalty to the input strides: given the choice between - // consecutive writes and consecutive reads, we would prefer consecutive - // writes. - constexpr double kPenalty = 1.01; - stride = std::min(a_stride * kPenalty, b_stride); - return std::make_tuple(l.contiguity, -stride); + return std::make_tuple(0, -a_stride, 0.0); + case ChunkContiguity::kNone: { + constexpr double kOutputPenalty = 1.01; + const double min_stride = + std::min(a_stride * kOutputPenalty, b_stride); + + // Pure tile-packing identity linearizers/delinearizers preserve the + // baseline std::min cost to keep 2D hardware tile pairings intact. + if (is_identity_permutation_) { + return std::make_tuple(l.contiguity, -min_stride, 0.0); + } + + // Model the memory traffic of placing this loop in an inner position + // based on the count of cache lines missed (bytes pulled from L2 to + // L1). Memory transfers occur at cache-line granularity (128-byte cache + // line pairs on modern CPUs like AMD Zen 3/4). + // - Strides < 128 bytes share cache lines across iterations (spatial + // reuse). + // - Strides >= 128 bytes incur a full cache-line miss per iteration. + constexpr double kCacheLineBytes = 128.0; + + const double misses_a = std::min(1.0, a_stride / kCacheLineBytes); + const double misses_b = std::min(1.0, b_stride / kCacheLineBytes); + const double base_cost = misses_a + misses_b * kOutputPenalty; + + // Break ties among loops with full cache-line misses using the minimum + // stride across both buffers (favoring loops with better locality in + // at least one buffer, such as tile interior loops over tile exterior + // loops). + return std::make_tuple(l.contiguity, -base_cost, -min_stride); + } } }; @@ -1417,6 +1456,7 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { } // Hard constraint 2: tile ordering. + // A tile interior MUST come after its corresponding tile exterior. if (l.tile_interior) { auto is_exterior = [&](const Loop& r) { return !r.tile_interior && r.dim_in_a == l.dim_in_a; @@ -1434,6 +1474,7 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { loop_order.push_back(std::move(remaining[best_idx])); remaining.erase(remaining.begin() + best_idx); } + VLOG(5) << "After loop ordering sort: " << absl::StrJoin(loop_order, ", ", [](std::string* out, const Loop& l) { diff --git a/third_party/xla/xla/pjrt/transpose.h b/third_party/xla/xla/pjrt/transpose.h index 9a91f53a488c76..1b7105dc9b1d60 100644 --- a/third_party/xla/xla/pjrt/transpose.h +++ b/third_party/xla/xla/pjrt/transpose.h @@ -353,6 +353,7 @@ class TransposePlan { absl::InlinedVector b_tiling_; bool a_is_tiled_ = false; bool b_is_tiled_ = false; + bool is_identity_permutation_ = false; // Per-chunk loop nests. Each loop nest has its own start/end bounds // representing one chunk of the work. diff --git a/third_party/xla/xla/python/ifrt/sharding.cc b/third_party/xla/xla/python/ifrt/sharding.cc index a641d51f29e315..eda18e355a077b 100644 --- a/third_party/xla/xla/python/ifrt/sharding.cc +++ b/third_party/xla/xla/python/ifrt/sharding.cc @@ -22,11 +22,9 @@ limitations under the License. #include #include #include -#include #include #include "absl/algorithm/container.h" -#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc index cb6ecf2d0d1857..139bba53c524ce 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.h" #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc index 298cd13872229a..b2e91574307647 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/python/pjrt_ifrt/pjrt_client.h" #include -#include #include #include #include diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc index b2502d0330e6da..3f3d42e2edbe33 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc index 74523156193a9f..6d0182b3ac60ea 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include #include -#include #include #include #include diff --git a/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc index e20a8d36e5c62e..68e5913ac29bbb 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc index 373f8508abcb46..6ffda12e9c988b 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc index 9f6593a954619f..02bc31872caf94 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 8ae6293d347548..0b0fd017b209df 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -661,93 +661,34 @@ xla_cc_test( cc_library( name = "riegeli_file_writer_factory", + srcs = ["riegeli_file_writer_factory.cc"], hdrs = ["riegeli_file_writer_factory.h"], + tags = if_google(["ignore_for_dep=third_party/riegeli/bytes/fd_writer.h"]), deps = [ - "@com_google_absl//absl/status", "@com_google_absl//absl/strings:string_view", "@riegeli//riegeli/bytes:writer", + "@tsl//tsl/platform", ] + if_google( - [":riegeli_file_writer_factory_google"], - [":riegeli_file_writer_factory_oss"], + ["@riegeli//riegeli/bytes:file_writer"], + ["@riegeli//riegeli/bytes:fd_writer"], ), ) -cc_library( - name = "riegeli_file_writer_factory_google", - srcs = [ - "riegeli_file_writer_factory.h", - "riegeli_file_writer_factory_google.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - # copybara:uncomment "@riegeli//riegeli/bytes:file_writer", - "@riegeli//riegeli/bytes:writer", - ], - alwayslink = 1, -) - -cc_library( - name = "riegeli_file_writer_factory_oss", - srcs = [ - "riegeli_file_writer_factory.h", - "riegeli_file_writer_factory_oss.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - "@riegeli//riegeli/bytes:fd_writer", - "@riegeli//riegeli/bytes:writer", - ], - alwayslink = 1, -) - cc_library( name = "riegeli_file_reader_factory", + srcs = ["riegeli_file_reader_factory.cc"], hdrs = ["riegeli_file_reader_factory.h"], + tags = if_google(["ignore_for_dep=third_party/riegeli/bytes/fd_reader.h"]), deps = [ "@com_google_absl//absl/strings:string_view", "@riegeli//riegeli/bytes:reader", + "@tsl//tsl/platform", ] + if_google( - [":riegeli_file_reader_factory_google"], - [":riegeli_file_reader_factory_oss"], + ["@riegeli//riegeli/bytes:file_reader"], + ["@riegeli//riegeli/bytes:fd_reader"], ), ) -cc_library( - name = "riegeli_file_reader_factory_google", - srcs = [ - "riegeli_file_reader_factory.h", - "riegeli_file_reader_factory_google.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - # copybara:uncomment "@riegeli//riegeli/bytes:file_reader", - "@riegeli//riegeli/bytes:reader", - ], - alwayslink = 1, -) - -cc_library( - name = "riegeli_file_reader_factory_oss", - srcs = [ - "riegeli_file_reader_factory.h", - "riegeli_file_reader_factory_oss.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - "@riegeli//riegeli/bytes:fd_reader", - "@riegeli//riegeli/bytes:reader", - ], - alwayslink = 1, -) - xla_cc_test( name = "dump_test", srcs = ["dump_test.cc"], diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc index 062cd41a9d6268..bdfef0a45c0dea 100644 --- a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc +++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc @@ -271,12 +271,17 @@ absl::StatusOr HloOpProfiler::MeasureOpChainDuration( std::minstd_rand0 engine; // Some operations have dynamic duration that depends on the input values. // Measure each operation with small and large inputs and average. - std::vector args_small = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/false) - .value(); - std::vector args_large = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/true) - .value(); + FakeArgumentsOptions small_options; + small_options.engine = &engine; + small_options.use_large_range = false; + ABSL_ASSIGN_OR_RETURN(std::vector args_small, + MakeFakeArguments(module.get(), small_options)); + + FakeArgumentsOptions large_options; + large_options.engine = &engine; + large_options.use_large_range = true; + ABSL_ASSIGN_OR_RETURN(std::vector args_large, + MakeFakeArguments(module.get(), large_options)); const absl::Time t_compile_start = absl::Now(); ABSL_ASSIGN_OR_RETURN(std::unique_ptr ex, runner_.CreateExecutable(std::move(module), diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index bcc343b0bab574..9cf9e39a46e95e 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -834,6 +834,7 @@ absl::Status LayoutAssignment::AddWhileConstraints( condition_layout.parameter_shape(0))); DCHECK(ShapeUtil::Compatible(body_layout.result_shape(), init->shape())); + bool is_copy_disabled_body = copy_disabled_while_computations_.contains(body); if (body_layout.result_layout() != body_layout.parameter_layout(0)) { VLOG(2) << "Reset %while body parameter layout: body=" << body->name() << " while=" << instruction->name() @@ -843,7 +844,19 @@ absl::Status LayoutAssignment::AddWhileConstraints( body_layout, current_priority_ + kNumberOfPropagationRounds, /*prop_result_layout=*/true, /*prop_parameter_layout=*/true); + if (is_copy_disabled_body) { + while_layout_changed_ = true; + } + } + if (is_copy_disabled_body) { + PropagateWhileLoopLayoutToSubcomputations( + body, body_layout.parameter_layout(0).shape(), + &body_layout.result_layout().shape(), + current_priority_ + kNumberOfPropagationRounds); } + + bool is_copy_disabled_condition = + copy_disabled_while_computations_.contains(condition); if (condition_layout.parameter_layout(0) != body_layout.parameter_layout(0)) { VLOG(2) << "Reset %while condition parameter layout: cond=" << condition->name() << " while=" << instruction->name() @@ -853,12 +866,387 @@ absl::Status LayoutAssignment::AddWhileConstraints( condition_constraint->ResetComputationLayout( condition_layout, current_priority_ + kNumberOfPropagationRounds, /*prop_result_layout=*/true, /*prop_parameter_layout=*/true); + if (is_copy_disabled_condition) { + while_layout_changed_ = true; + } + } + if (is_copy_disabled_condition) { + PropagateWhileLoopLayoutToSubcomputations( + condition, condition_layout.parameter_layout(0).shape(), + &condition_layout.result_layout().shape(), + current_priority_ + kNumberOfPropagationRounds); } ABSL_RETURN_IF_ERROR(SetOperandLayout(body_layout.result_shape(), instruction, 0)); return SetInstructionLayout(body_layout.result_shape(), instruction); } +void LayoutAssignment::PropagateWhileLoopLayoutToSubcomputations( + HloComputation* computation, const Shape& param_shape, + const Shape* result_shape, int64_t priority) { + if (computation == nullptr || computation->num_parameters() == 0) { + return; + } + + auto HasAnyLayout = [](const Shape& shape) -> bool { + bool has = false; + ShapeUtil::ForEachSubshape( + shape, [&has](const Shape& subshape, const ShapeIndex&) { + if (subshape.has_layout()) { + has = true; + } + }); + return has; + }; + + absl::flat_hash_map instruction_layouts; + if (HasAnyLayout(param_shape) && + ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(), + param_shape)) { + instruction_layouts[computation->parameter_instruction(0)] = param_shape; + } + + // Forward propagation from parameter through the computation. + for (HloInstruction* instr : computation->MakeInstructionPostOrder()) { + if (instr == computation->parameter_instruction(0)) { + continue; + } + switch (instr->opcode()) { + case HloOpcode::kGetTupleElement: { + auto it = instruction_layouts.find(instr->operand(0)); + if (it != instruction_layouts.end() && it->second.IsTuple() && + instr->tuple_index() < ShapeUtil::TupleElementCount(it->second)) { + Shape subshape = + ShapeUtil::GetTupleElementShape(it->second, instr->tuple_index()); + if (ShapeUtil::Compatible(instr->shape(), subshape)) { + instruction_layouts[instr] = std::move(subshape); + } + } + break; + } + case HloOpcode::kTuple: { + std::vector subshapes; + subshapes.reserve(instr->operand_count()); + bool any_known = false; + for (int64_t i = 0; i < instr->operand_count(); ++i) { + auto it = instruction_layouts.find(instr->operand(i)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->operand(i)->shape(), it->second)) { + subshapes.push_back(it->second); + any_known = true; + } else { + Shape cleared = instr->operand(i)->shape(); + LayoutUtil::ClearLayout(&cleared); + subshapes.push_back(std::move(cleared)); + } + } + if (any_known) { + Shape tuple_shape = ShapeUtil::MakeTupleShape(subshapes); + if (ShapeUtil::Compatible(instr->shape(), tuple_shape)) { + instruction_layouts[instr] = std::move(tuple_shape); + } + } + break; + } + case HloOpcode::kSelect: { + Shape select_shape; + bool has_select_shape = false; + auto it_true = instruction_layouts.find(instr->operand(1)); + auto it_false = instruction_layouts.find(instr->operand(2)); + if (it_true != instruction_layouts.end() && + HasAnyLayout(it_true->second) && + ShapeUtil::Compatible(instr->shape(), it_true->second)) { + select_shape = it_true->second; + has_select_shape = true; + } else if (it_false != instruction_layouts.end() && + HasAnyLayout(it_false->second) && + ShapeUtil::Compatible(instr->shape(), it_false->second)) { + select_shape = it_false->second; + has_select_shape = true; + } + if (has_select_shape) { + instruction_layouts[instr] = std::move(select_shape); + } + break; + } + case HloOpcode::kCopy: { + auto it = instruction_layouts.find(instr->operand(0)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->shape(), it->second)) { + Shape copy_shape = it->second; + instruction_layouts[instr] = std::move(copy_shape); + } + break; + } + case HloOpcode::kConditional: { + Shape branch_result_shape; + bool has_branch_result_shape = false; + for (int j = 0; j < instr->branch_count(); ++j) { + HloInstruction* branch_arg = instr->mutable_operand(j + 1); + HloComputation* branch_comp = instr->branch_computation(j); + auto it = instruction_layouts.find(branch_arg); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + !ShapeUtil::IsEmptyTuple(it->second) && + ShapeUtil::Compatible(branch_arg->shape(), it->second)) { + PropagateWhileLoopLayoutToSubcomputations( + branch_comp, it->second, /*result_shape=*/nullptr, priority); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr && + branch_constraints->computation_constraint() + .result_layout_is_set()) { + const Shape& res = branch_constraints->computation_layout() + .result_layout() + .shape(); + if (ShapeUtil::Compatible(instr->shape(), res)) { + branch_result_shape = res; + has_branch_result_shape = true; + } + } + } + } + if (!has_branch_result_shape) { + auto it = instruction_layouts.find(instr); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->shape(), it->second)) { + branch_result_shape = it->second; + has_branch_result_shape = true; + } + } + if (has_branch_result_shape) { + instruction_layouts[instr] = branch_result_shape; + for (int j = 0; j < instr->branch_count(); ++j) { + HloComputation* branch_comp = instr->branch_computation(j); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr) { + ComputationLayout branch_layout = + branch_constraints->computation_layout(); + bool prop_branch_result = false; + ShapeUtil::ForEachSubshape( + branch_result_shape, + [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + branch_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_branch_result = true; + } + }); + if (prop_branch_result) { + branch_constraints->mutable_computation_constraint() + ->ResetComputationLayout( + branch_layout, priority, + /*prop_result_layout=*/true, + /*prop_parameter_layout=*/ + branch_constraints->computation_constraint() + .parameter_layout_is_set()); + } + } + } + } + break; + } + case HloOpcode::kWhile: { + HloInstruction* init = instr->mutable_operand(0); + auto it = instruction_layouts.find(init); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(init->shape(), it->second)) { + PropagateWhileLoopLayoutToSubcomputations( + instr->while_body(), it->second, &it->second, priority); + PropagateWhileLoopLayoutToSubcomputations( + instr->while_condition(), it->second, /*result_shape=*/nullptr, + priority); + } + break; + } + case HloOpcode::kCall: { + HloComputation* callee = instr->to_apply(); + if (callee != nullptr && + callee->num_parameters() == instr->operand_count()) { + auto* callee_constraints = mutable_computation_constraints(callee); + if (callee_constraints != nullptr) { + ComputationLayout callee_layout = + callee_constraints->computation_layout(); + bool any_param_set = false; + for (int64_t i = 0; i < instr->operand_count(); ++i) { + auto it = instruction_layouts.find(instr->operand(i)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->operand(i)->shape(), + it->second)) { + *callee_layout.mutable_parameter_layout(i) = + ShapeLayout(it->second); + any_param_set = true; + } + } + if (any_param_set) { + callee_constraints->mutable_computation_constraint() + ->ResetComputationLayout(callee_layout, priority, + /*prop_result_layout=*/false, + /*prop_parameter_layout=*/true); + } + } + } + break; + } + default: + break; + } + } + + // Backward propagation from result_shape if provided. + if (result_shape != nullptr && HasAnyLayout(*result_shape) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + *result_shape)) { + absl::flat_hash_map root_layouts; + root_layouts[computation->root_instruction()] = *result_shape; + auto post_order = computation->MakeInstructionPostOrder(); + for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) { + HloInstruction* instr = *it; + auto root_it = root_layouts.find(instr); + if (root_it == root_layouts.end() || !HasAnyLayout(root_it->second)) { + continue; + } + Shape current_layout_shape = root_it->second; + switch (instr->opcode()) { + case HloOpcode::kTuple: { + if (current_layout_shape.IsTuple()) { + for (int64_t i = 0; i < instr->operand_count(); ++i) { + if (i < ShapeUtil::TupleElementCount(current_layout_shape)) { + Shape elem_shape = + ShapeUtil::GetTupleElementShape(current_layout_shape, i); + if (ShapeUtil::Compatible(instr->operand(i)->shape(), + elem_shape)) { + root_layouts[instr->operand(i)] = std::move(elem_shape); + } + } + } + } + break; + } + case HloOpcode::kSelect: { + if (ShapeUtil::Compatible(instr->operand(1)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(1)] = current_layout_shape; + } + if (ShapeUtil::Compatible(instr->operand(2)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(2)] = current_layout_shape; + } + break; + } + case HloOpcode::kCopy: { + if (ShapeUtil::Compatible(instr->operand(0)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(0)] = current_layout_shape; + } + break; + } + case HloOpcode::kGetTupleElement: { + HloInstruction* gte_op = instr->mutable_operand(0); + if (gte_op->opcode() == HloOpcode::kConditional && + instr->tuple_index() < + ShapeUtil::TupleElementCount(gte_op->shape())) { + auto cond_it = root_layouts.find(gte_op); + if (cond_it == root_layouts.end()) { + root_layouts[gte_op] = gte_op->shape(); + } + if (root_layouts[gte_op].IsTuple()) { + *ShapeUtil::GetMutableSubshape(&root_layouts[gte_op], + {instr->tuple_index()}) = + current_layout_shape; + } + } + break; + } + case HloOpcode::kConditional: { + if (ShapeUtil::Compatible(instr->shape(), current_layout_shape)) { + for (int j = 0; j < instr->branch_count(); ++j) { + HloComputation* branch_comp = instr->branch_computation(j); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr) { + ComputationLayout branch_layout = + branch_constraints->computation_layout(); + bool prop_branch_result = false; + ShapeUtil::ForEachSubshape( + current_layout_shape, + [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + branch_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_branch_result = true; + } + }); + if (prop_branch_result) { + branch_constraints->mutable_computation_constraint() + ->ResetComputationLayout( + branch_layout, priority, + /*prop_result_layout=*/true, + /*prop_parameter_layout=*/ + branch_constraints->computation_constraint() + .parameter_layout_is_set()); + } + } + } + } + break; + } + default: + break; + } + } + } + + // Update this computation's own ComputationLayout. + auto* comp_constraints = mutable_computation_constraints(computation); + if (comp_constraints != nullptr) { + ComputationLayout comp_layout = comp_constraints->computation_layout(); + bool prop_param = false; + if (HasAnyLayout(param_shape) && + ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(), + param_shape)) { + ShapeUtil::ForEachSubshape( + param_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_parameter_layout(0)->ResetLayout( + subshape.layout(), index); + prop_param = true; + } + }); + } + bool prop_result = false; + auto it = instruction_layouts.find(computation->root_instruction()); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + it->second)) { + ShapeUtil::ForEachSubshape( + it->second, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_result = true; + } + }); + } else if (result_shape != nullptr && HasAnyLayout(*result_shape) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + *result_shape)) { + ShapeUtil::ForEachSubshape( + *result_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_result = true; + } + }); + } + if (prop_param || prop_result) { + comp_constraints->mutable_computation_constraint() + ->ResetComputationLayout(comp_layout, priority, prop_result, + prop_param); + } + } +} + absl::Status LayoutAssignment::AddConditionalConstraints( HloInstruction* instruction) { if (computation_layouts_.find(instruction->branch_computation(0)) == @@ -1262,12 +1650,24 @@ absl::Status CheckWhileLayout( auto init_shape = while_inst->operand(0)->shape(); TF_RET_CHECK( condition_computation_layout.parameter_layout(0).MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " cond param layout=" + << condition_computation_layout.parameter_layout(0).ToString() + << " init_shape=" << init_shape.ToString(); TF_RET_CHECK(body_computation_layout.parameter_layout(0).MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " body param layout=" + << body_computation_layout.parameter_layout(0).ToString() + << " init_shape=" << init_shape.ToString(); TF_RET_CHECK(body_computation_layout.result_layout().MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); - TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape())); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " body result layout=" + << body_computation_layout.result_layout().ToString() + << " init_shape=" << init_shape.ToString(); + TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape())) + << " while_inst=" << while_inst->name() + << " while shape=" << while_inst->shape().ToString() + << " init_shape=" << init_shape.ToString(); return absl::OkStatus(); } @@ -1476,7 +1876,8 @@ absl::Status LayoutAssignment::CopyOperandIfLayoutsDiffer( VLOG(2) << "Operand " << operand->ToString() << " layout does not match " << operand_layout.ToString() << " in " << instruction->ToString(); - if (IsWhileLoopCopyDisabled(*instruction)) { + if (IsWhileLoopCopyDisabled(*instruction) && + !instruction->parent()->caller_instructions(HloOpcode::kWhile).empty()) { HloComputation* comp = instruction->parent(); HloInstruction* param = comp->parameter_instruction(0); ShapeIndex index = {}; @@ -2682,7 +3083,8 @@ absl::Status LayoutAssignment::CalculateComputationLayout( callee->mutable_computation_constraint(); ComputationLayout callee_layout = callee_constraint->computation_layout(); if (callee_constraint->priority() < priority || - conditional_mismatch_.count(callee->computation()) > 0) { + conditional_mismatch_.count(callee->computation()) > 0 || + copy_disabled_while_computations_.contains(callee->computation())) { if (conditional_mismatch_.count(callee->computation()) == 0 && UpdateLayout(result, callee_layout.mutable_result_layout())) { VLOG(2) << "Setting result layout from : " << result->ToString() @@ -2762,7 +3164,8 @@ absl::Status LayoutAssignment::CalculateComputationLayout( } // Reset the layout of the current computation from its body. if (current_priority_ == 0 || - conditional_mismatch_.count(constraints->computation()) > 0) { + conditional_mismatch_.count(constraints->computation()) > 0 || + copy_disabled_while_computations_.contains(constraints->computation())) { ABSL_RETURN_IF_ERROR(SetCalleeLayout( constraints->computation()->root_instruction(), constraints->computation()->parameter_instructions(), constraints, @@ -3240,8 +3643,11 @@ absl::StatusOr LayoutAssignment::RunImpl( // Layouts are allowed to flow naturally in the first round, and any detected // inconsistencies at boundary instructions are resolved with higher-priority // constraints in subsequent rounds. - for (int64_t i = 0; changed || i < kNumberOfPropagationRounds; ++i) { + for (int64_t i = 0; + (changed || i < kNumberOfPropagationRounds) && i < kMaxPropagationRounds; + ++i) { changed = false; + while_layout_changed_ = false; VLOG(1) << "Running " << (i == 0 ? "un" : "") << "constrained pass"; ABSL_RETURN_IF_ERROR(ClearPreviousPassSideEffects(module, execution_threads)); // Layouts are propagated within each computation. In the first round, @@ -3260,6 +3666,7 @@ absl::StatusOr LayoutAssignment::RunImpl( ABSL_ASSIGN_OR_RETURN(bool aliasing_changed, ResolveInputOutputAliasing(module, entry_constraint)); changed |= aliasing_changed; + changed |= while_layout_changed_; } // All logical buffers should have constraints at this point. All that @@ -3450,6 +3857,7 @@ absl::Status LayoutAssignment::Init(HloModule* module) { computation_layouts_.clear(); conditional_mismatch_.clear(); current_priority_ = LayoutConstraint::kBeginningPriority; + while_layout_changed_ = false; // Clear all the copies which have been added, and all the related // instructions (like GTE and tuples). if (!added_copies_.empty()) { diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h index 9c4481fe2a4f4a..daf52cdc9c0c54 100644 --- a/third_party/xla/xla/service/layout_assignment.h +++ b/third_party/xla/xla/service/layout_assignment.h @@ -598,6 +598,7 @@ class LayoutAssignment : public HloModulePass { ComputationLayout& saved_entry_computation_layout() { return saved_entry_computation_layout_; } + virtual bool NegotiateLayout(const HloInstruction* instruction, const Layout& new_layout, const Layout& existing_layout, @@ -803,6 +804,13 @@ class LayoutAssignment : public HloModulePass { absl::Status AddAsyncDoneConstraints(HloInstruction* instruction, LayoutConstraints* constraints); + // Propagates while loop parameter and result layouts to subcomputations (such + // as conditionals) within the while body or condition. + void PropagateWhileLoopLayoutToSubcomputations(HloComputation* computation, + const Shape& param_shape, + const Shape* result_shape, + int64_t priority); + // Propagates layout constraints from the caller instruction into the inner // async sub-computation. // This is the forward propagation step: it takes the layouts of the operands @@ -952,6 +960,7 @@ class LayoutAssignment : public HloModulePass { protected: static constexpr int64_t kNumberOfPropagationRounds = 2; + static constexpr int64_t kMaxPropagationRounds = 6; // Sets up the copy instruction according to the characteristic (sharding, // metadata, ...) of the reference instruction. The index argument is used // when the instruction is a tuple, and in such case the index represents @@ -1072,6 +1081,10 @@ class LayoutAssignment : public HloModulePass { // Stores the set of while computations that have copy disabled. absl::flat_hash_set copy_disabled_while_computations_; + + // Tracks whether while loop parameter/condition layouts changed in the + // current propagation round and require another round to converge. + bool while_layout_changed_ = false; }; } // namespace xla diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc index d0a9091ead93db..659b95b064e256 100644 --- a/third_party/xla/xla/service/layout_assignment_test.cc +++ b/third_party/xla/xla/service/layout_assignment_test.cc @@ -2311,6 +2311,84 @@ ENTRY main { ExpectLayoutIs(ShapeUtil::GetSubshape(body_param_shape, {0}), {0, 1}); } +TEST_F(LayoutAssignmentTest, + RespectsDisableWhileLoopCopiesWithConditionalSubcomputation) { + const char* module_str = R"( +HloModule t + +while_condition { + tuple = (s32[2,8]{1,0}, u32[], pred[]) parameter(0) + i = u32[] get-tuple-element(tuple), index=1 + n = u32[] constant(8) + ROOT predicate = pred[] compare(i, n), direction=LT +} + +branch_true { + arg = s32[2,8]{1,0} parameter(0) + ROOT custom = s32[2,8]{0,1} custom-call(arg), custom_call_target="baz", + operand_layout_constraints={s32[2,8]{0,1}} +} + +branch_false { + ROOT arg = s32[2,8]{1,0} parameter(0) +} + +while_body { + tuple = (s32[2,8]{1,0}, u32[], pred[]) parameter(0) + input = s32[2,8]{1,0} get-tuple-element(tuple), index=0 + i = u32[] get-tuple-element(tuple), index=1 + p = pred[] get-tuple-element(tuple), index=2 + c1 = u32[] constant(1) + i_ = add(i, c1) + cond = s32[2,8]{0,1} conditional(p, input, input), true_computation=branch_true, false_computation=branch_false + ROOT tuple1 = (s32[2,8]{0,1}, u32[], pred[]) tuple(cond, i_, p) +} + +ENTRY main { + input = s32[2,8]{1,0} parameter(0) + c0 = u32[] constant(0) + p0 = pred[] constant(true) + tuple = (s32[2,8]{1,0}, u32[], pred[]) tuple(input, c0, p0) + tuple_ = (s32[2,8]{1,0}, u32[], pred[]) while(tuple), condition=while_condition, body=while_body, frontend_attributes={xla_disable_while_loop_copies="true"} + ROOT output_ = s32[2,8]{1,0} get-tuple-element(tuple_), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); + EXPECT_OK(RunLayoutAssignmentPass(m.get())); + + HloComputation* body = m->GetComputationWithName("while_body"); + ASSERT_NE(body, nullptr); + for (HloInstruction* inst : body->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + + HloComputation* branch_true = m->GetComputationWithName("branch_true"); + ASSERT_NE(branch_true, nullptr); + for (HloInstruction* inst : branch_true->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + ExpectLayoutIs(branch_true->parameter_instruction(0)->shape(), {0, 1}); + ExpectLayoutIs(branch_true->root_instruction()->shape(), {0, 1}); + + HloComputation* branch_false = m->GetComputationWithName("branch_false"); + ASSERT_NE(branch_false, nullptr); + for (HloInstruction* inst : branch_false->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + ExpectLayoutIs(branch_false->parameter_instruction(0)->shape(), {0, 1}); + ExpectLayoutIs(branch_false->root_instruction()->shape(), {0, 1}); + + HloInstruction* cond = FindInstruction(m.get(), "cond"); + ASSERT_NE(cond, nullptr); + ExpectLayoutIs(cond->shape(), {0, 1}); + + const Shape& body_param_shape = body->parameter_instruction(0)->shape(); + ExpectLayoutIs(ShapeUtil::GetSubshape(body_param_shape, {0}), {0, 1}); + ExpectLayoutIs(ShapeUtil::GetSubshape(body->root_instruction()->shape(), {0}), + {0, 1}); +} + TEST_F(LayoutAssignmentTest, HloBufferLayoutUnconstrained) { const char* module_str = R"( HloModule test diff --git a/third_party/xla/xla/service/riegeli_file_reader_factory_google.cc b/third_party/xla/xla/service/riegeli_file_reader_factory.cc similarity index 84% rename from third_party/xla/xla/service/riegeli_file_reader_factory_google.cc rename to third_party/xla/xla/service/riegeli_file_reader_factory.cc index 1f60b9689307d5..02f2f9910afa6e 100644 --- a/third_party/xla/xla/service/riegeli_file_reader_factory_google.cc +++ b/third_party/xla/xla/service/riegeli_file_reader_factory.cc @@ -13,18 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/service/riegeli_file_reader_factory.h" + #include #include "absl/strings/string_view.h" -#include "riegeli/bytes/file_reader.h" #include "riegeli/bytes/reader.h" -#include "xla/service/riegeli_file_reader_factory.h" +#include "tsl/platform/platform.h" + +#if TSL_IS_IN_OSS +#include "riegeli/bytes/fd_reader.h" +#else +#include "riegeli/bytes/file_reader.h" +#endif namespace xla { std::unique_ptr CreateRiegeliFileReader( absl::string_view filename) { +#if TSL_IS_IN_OSS + return std::make_unique>(filename); +#else return std::make_unique>(filename); +#endif } } // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc b/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc deleted file mode 100644 index 479bc0939f2473..00000000000000 --- a/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc +++ /dev/null @@ -1,30 +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 - -#include "absl/strings/string_view.h" -#include "riegeli/bytes/fd_reader.h" -#include "riegeli/bytes/reader.h" -#include "xla/service/riegeli_file_reader_factory.h" - -namespace xla { - -std::unique_ptr CreateRiegeliFileReader( - absl::string_view filename) { - return std::make_unique>(filename); -} - -} // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_writer_factory_google.cc b/third_party/xla/xla/service/riegeli_file_writer_factory.cc similarity index 84% rename from third_party/xla/xla/service/riegeli_file_writer_factory_google.cc rename to third_party/xla/xla/service/riegeli_file_writer_factory.cc index c236236b66d140..4b72cf47b0507d 100644 --- a/third_party/xla/xla/service/riegeli_file_writer_factory_google.cc +++ b/third_party/xla/xla/service/riegeli_file_writer_factory.cc @@ -13,18 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/service/riegeli_file_writer_factory.h" + #include #include "absl/strings/string_view.h" -#include "riegeli/bytes/file_writer.h" #include "riegeli/bytes/writer.h" -#include "xla/service/riegeli_file_writer_factory.h" +#include "tsl/platform/platform.h" + +#if TSL_IS_IN_OSS +#include "riegeli/bytes/fd_writer.h" +#else +#include "riegeli/bytes/file_writer.h" +#endif namespace xla { std::unique_ptr CreateRiegeliFileWriter( absl::string_view filename) { +#if TSL_IS_IN_OSS + return std::make_unique>(filename); +#else return std::make_unique>(filename); +#endif } } // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc b/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc deleted file mode 100644 index 19fb7e9297f897..00000000000000 --- a/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc +++ /dev/null @@ -1,30 +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 - -#include "absl/strings/string_view.h" -#include "riegeli/bytes/fd_writer.h" -#include "riegeli/bytes/writer.h" -#include "xla/service/riegeli_file_writer_factory.h" - -namespace xla { - -std::unique_ptr CreateRiegeliFileWriter( - absl::string_view filename) { - return std::make_unique>(filename); -} - -} // namespace xla diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index ac980eb7a61ec1..bb0b2f2448d70c 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -440,7 +440,8 @@ void ConstraintPropagator::ComputeMaxAddReductionElementsPerExp( consumer_add_reduction_elements = it->second; } - if (instruction->opcode() == HloOpcode::kExp) { + if (instruction->opcode() == HloOpcode::kExp || + instruction->opcode() == HloOpcode::kExpm1) { max_add_reduction_elements_per_exp_[instruction] = consumer_add_reduction_elements; } @@ -524,7 +525,8 @@ absl::Status ConstraintPropagator::SeedConstraints( // Output is guaranteed to be non-negative. states_[inst].AddConstraint(ConstraintInterval::Positive()); break; - case HloOpcode::kExp: { + case HloOpcode::kExp: + case HloOpcode::kExpm1: { // Safe domain [-max_log, max_log] prevents floating point overflow. int64_t reduction_elements = GetMaxAddReductionElementsForExp(inst); double max_log = @@ -1599,6 +1601,33 @@ void ConstraintPropagator::PropagateExpApprox( ConstraintInterval{x_min, x_max, /*exclude_zero=*/false}); } +void ConstraintPropagator::PropagateExpm1Approx( + const HloInstruction* instruction, + const ConstraintInterval& output_interval) { + // For Y = expm1(X) = exp(X) - 1 with Y in [y_min, y_max]: + // Since expm1(X) is monotonically strictly increasing on real numbers: + // y_min <= expm1(X) <= y_max <=> ln(y_min + 1) <= X <= ln(y_max + 1). + // + // Since expm1(X) > -1 for all real X: + // If y_max <= -1.0, expm1(X) <= y_max is impossible for real numbers. + if (output_interval.max <= -1.0) { + states_[instruction->operand(0)].AddConstraint( + ConstraintInterval{1.0, -1.0, /*exclude_zero=*/false}); + return; + } + + double x_min = output_interval.min > -1.0 ? std::log1p(output_interval.min) + : ConstraintInterval::kMin; + double x_max = output_interval.max < ConstraintInterval::kMax && + output_interval.max > -1.0 + ? std::log1p(output_interval.max) + : ConstraintInterval::kMax; + bool exclude_zero = output_interval.exclude_zero && + output_interval.min <= 0.0 && output_interval.max >= 0.0; + states_[instruction->operand(0)].AddConstraint( + ConstraintInterval{x_min, x_max, exclude_zero}); +} + void ConstraintPropagator::PropagatePowerApprox( const HloInstruction* instruction, const ConstraintInterval& output_interval) { @@ -1679,6 +1708,9 @@ absl::Status ConstraintPropagator::PropagateConstraintsApprox( case HloOpcode::kExp: PropagateExpApprox(instruction, output_interval); break; + case HloOpcode::kExpm1: + PropagateExpm1Approx(instruction, output_interval); + break; case HloOpcode::kPower: PropagatePowerApprox(instruction, output_interval); break; diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h index 9fa12d6fb24c36..22ff086bd5b2b8 100644 --- a/third_party/xla/xla/tests/constraint_propagator.h +++ b/third_party/xla/xla/tests/constraint_propagator.h @@ -146,6 +146,8 @@ class ConstraintPropagator { const ConstraintInterval& output_interval); void PropagateExpApprox(const HloInstruction* instruction, const ConstraintInterval& output_interval); + void PropagateExpm1Approx(const HloInstruction* instruction, + const ConstraintInterval& output_interval); void PropagatePowerApprox(const HloInstruction* instruction, const ConstraintInterval& output_interval); diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 47d0b515bc739e..b216d3743224c1 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -1503,5 +1503,83 @@ ENTRY main { *module->GetComputationWithName("min_computation")), IdentityElementType::kMaximum); } + +TEST_F(ConstraintPropagatorTest, Expm1SeedConstrainsSafeDomain) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = f32[8,128] parameter(0) + ROOT root = f32[8,128] exponential-minus-one(x) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + // Safe domain [-max_log, max_log] = [-4.0, 4.0] for F32 prevents overflow. + EXPECT_DOUBLE_EQ(x_int.min, -4.0); + EXPECT_DOUBLE_EQ(x_int.max, 4.0); + EXPECT_FALSE(x_int.exclude_zero); +} + +TEST_F(ConstraintPropagatorTest, Expm1BackwardPropagationStrictPositive) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = f32[8,128] parameter(0) + expm1 = f32[8,128] exponential-minus-one(x) + ROOT log = f32[8,128] log(expm1) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + // log(expm1(x)) requires expm1(x) > 0. + // Since expm1(x) > 0 <=> x > log1p(0) = 0, x must be strictly positive. + EXPECT_TRUE(x_int.IsPositiveStrict()); + EXPECT_DOUBLE_EQ(x_int.min, 0.0); +} + +TEST_F(ConstraintPropagatorTest, NanToNumLogFusionConstraints) { + constexpr absl::string_view kHloString = R"hlo( +HloModule fusion.30343, is_scheduled=true + +%fused_computation (param_0: f32[2,6], param_1: f32[2,6]) -> f32[2,6] { + %param_0 = f32[2,6] parameter(0) + %param_1 = f32[2,6] parameter(1) + %expm1 = f32[2,6] exponential-minus-one(%param_0) + %neg = f32[2,6] negate(%expm1) + %add = f32[2,6] add(%param_1, %neg) + %div = f32[2,6] divide(%param_1, %add) + %ne = pred[2,6] compare(%div, %div), direction=NE + %zero = f32[] constant(0) + %broadcast_zero = f32[2,6] broadcast(%zero), dimensions={} + %select_nan = f32[2,6] select(%ne, %broadcast_zero, %div) + ROOT %log = f32[2,6] log(%select_nan) +} + +ENTRY %fusion.30343 (parameter.0: f32[2,6], parameter.1: f32[2,6]) -> f32[2,6] { + %parameter.0 = f32[2,6] parameter(0) + %parameter.1 = f32[2,6] parameter(1) + ROOT %fusion.30343 = f32[2,6] fusion( + %parameter.0, %parameter.1), kind=kLoop, calls=%fused_computation +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto p0_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + auto p1_int = states[module->entry_computation()->parameter_instruction(1)] + .GetConstraintInterval(); + + EXPECT_TRUE(p0_int.IsNegativeStrict()); + EXPECT_DOUBLE_EQ(p0_int.min, -4.0); + EXPECT_TRUE(p1_int.IsPositiveStrict()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h b/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h index 98ff2277fa4704..48ea1e3061bbb7 100644 --- a/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h +++ b/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h @@ -127,11 +127,10 @@ class HloRunnerAgnosticReferenceMixin : public T { const std::function& reference_preprocessor = nullptr, const std::function& test_preprocessor = nullptr, const std::optional args_max_bits_of_precision = std::nullopt) { + FakeArgumentsOptions options; + options.max_bits_of_precision = args_max_bits_of_precision; const absl::StatusOr> fake_arguments = - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - args_max_bits_of_precision); + MakeFakeArguments(module.get(), options); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status().message(); } diff --git a/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc b/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc index 48195abd4cbcb9..ccb91b2b5bc2ff 100644 --- a/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc +++ b/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc @@ -295,11 +295,8 @@ HloRunnerAgnosticTestBase::RunAndCompareTwoModulesReplicated( std::unique_ptr module_0, std::unique_ptr module_1, const bool run_hlo_passes, const bool use_threads, const std::optional& error) { - const absl::StatusOr> fake_arguments = MakeFakeArguments( - /*module=*/module_0.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt); + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.get()); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -372,9 +369,10 @@ ::testing::AssertionResult HloRunnerAgnosticTestBase::RunAndCompareTwoModules( << absl::StrJoin(mismatches, ", "); } - const absl::StatusOr> fake_arguments = MakeFakeArguments( - module_0.get(), /*pseudo_random=*/true, /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, args_max_bits_of_precision); + FakeArgumentsOptions options; + options.max_bits_of_precision = args_max_bits_of_precision; + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.get(), options); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -483,11 +481,8 @@ HloRunnerAgnosticTestBase::RunAndCompareTwoExecutables( << "Error : mismatching parameter shapes for parameters " << absl::StrJoin(mismatches, ", "); } - absl::StatusOr> fake_arguments = MakeFakeArguments( - /*module=*/module_0.value(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt); + absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.value()); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -516,11 +511,18 @@ ::testing::AssertionResult HloRunnerAgnosticTestBase::Run( << "Error while parsing HLO text format: " << module.status().ToString(); } - const std::vector fake_arguments = - MakeFakeArguments(module->get(), use_random_data).value(); + FakeArgumentsOptions options; + options.pseudo_random = use_random_data; + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module->get(), options); + if (!fake_arguments.ok()) { + return ::testing::AssertionFailure() + << "Error while generating fake arguments: " + << fake_arguments.status().ToString(); + } std::vector fake_argument_ptrs; absl::c_transform( - fake_arguments, std::back_inserter(fake_argument_ptrs), + *fake_arguments, std::back_inserter(fake_argument_ptrs), [](const Literal& literal) { return const_cast(&literal); }); if (backend_config) { diff --git a/third_party/xla/xla/tests/test_utils.cc b/third_party/xla/xla/tests/test_utils.cc index a34c66b735cb38..c99759665732fd 100644 --- a/third_party/xla/xla/tests/test_utils.cc +++ b/third_party/xla/xla/tests/test_utils.cc @@ -338,36 +338,16 @@ absl::StatusOr MakeConstrainedArgument( } // namespace absl::StatusOr> MakeFakeArguments( - const HloModule* module, bool pseudo_random, bool use_large_range, - bool treat_gte_as_data_formatting, - std::optional max_bits_of_precision, std::minstd_rand0* engine, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { - if (!pseudo_random) { - return MakeFakeArguments(module, nullptr, use_large_range, - treat_gte_as_data_formatting, - max_bits_of_precision, generate_aligned_ds_indices, - get_index_known_zeroes); - } - if (engine == nullptr) { - auto new_engine = - pseudo_random ? std::make_unique() : nullptr; - return MakeFakeArguments(module, new_engine.get(), use_large_range, - treat_gte_as_data_formatting, - max_bits_of_precision, generate_aligned_ds_indices, - get_index_known_zeroes); + const HloModule* module, const FakeArgumentsOptions& options) { + std::unique_ptr default_engine; + std::minstd_rand0* engine = options.engine; + if (!options.pseudo_random) { + engine = nullptr; + } else if (engine == nullptr) { + default_engine = std::make_unique(); + engine = default_engine.get(); } - return MakeFakeArguments(module, engine, use_large_range, - treat_gte_as_data_formatting, max_bits_of_precision, - generate_aligned_ds_indices, get_index_known_zeroes); -} -absl::StatusOr> MakeFakeArguments( - const HloModule* module, std::minstd_rand0* engine, bool use_large_range, - bool treat_gte_as_data_formatting, - std::optional max_bits_of_precision, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { ABSL_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(*module)); const auto params = module->entry_computation()->parameter_instructions(); std::vector arguments(params.size()); @@ -386,33 +366,35 @@ absl::StatusOr> MakeFakeArguments( ABSL_ASSIGN_OR_RETURN( arguments[i], MakeConstrainedArgument( - *dataflow, *params[i], param_shape, engine, use_large_range, - treat_gte_as_data_formatting, max_bits_of_precision, - generate_aligned_ds_indices, get_index_known_zeroes)); + *dataflow, *params[i], param_shape, engine, options.use_large_range, + options.treat_gte_as_data_formatting, options.max_bits_of_precision, + options.generate_aligned_ds_indices, + options.get_index_known_zeroes)); } return std::move(arguments); } absl::StatusOr> MakeDataflowConstrainedArguments( - const HloModule* module, std::minstd_rand0* engine, bool use_large_range, - std::optional max_bits_of_precision, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { + const HloModule* module, const FakeArgumentsOptions& options) { std::unique_ptr default_engine; - if (engine == nullptr) { + std::minstd_rand0* engine = options.engine; + if (!options.pseudo_random) { + engine = nullptr; + } else if (engine == nullptr) { default_engine = std::make_unique(); engine = default_engine.get(); } - ABSL_ASSIGN_OR_RETURN(auto constraint_states, - ConstraintPropagator::Run(*module, get_index_known_zeroes)); + ABSL_ASSIGN_OR_RETURN( + auto constraint_states, + ConstraintPropagator::Run(*module, options.get_index_known_zeroes)); auto make_literal_for_state = [&](const Shape& shape, const ConstraintState& state, absl::string_view target_name) -> absl::StatusOr { ConstraintInterval interval = state.GetConstraintInterval(); StructuralConstraints structure = state.GetStructuralConstraints(); - if (!generate_aligned_ds_indices) { + if (!options.generate_aligned_ds_indices) { structure.alignment = std::nullopt; } std::optional> limit = std::nullopt; @@ -454,8 +436,8 @@ absl::StatusOr> MakeDataflowConstrainedArguments( limit = {min_val, max_val}; } return MakeFakeLiteral(shape, engine, limit, structure.needs_sorted_indices, - structure.no_duplicates, use_large_range, - max_bits_of_precision, structure.alignment, + structure.no_duplicates, options.use_large_range, + options.max_bits_of_precision, structure.alignment, structure.known_zeroes_mask, /*float_generator=*/nullptr, interval); }; diff --git a/third_party/xla/xla/tests/test_utils.h b/third_party/xla/xla/tests/test_utils.h index 23976222b8a562..da152feeb74adf 100644 --- a/third_party/xla/xla/tests/test_utils.h +++ b/third_party/xla/xla/tests/test_utils.h @@ -65,6 +65,53 @@ class PseudorandomGenerator { using GetIndexKnownZeroesFn = std::function(const HloInstruction*, int64_t)>; +// Options for generating fake arguments with MakeFakeArguments and +// MakeDataflowConstrainedArguments. +struct FakeArgumentsOptions { + // Optional random number generator. Passing a generator enables generation + // of different random values across sequential calls by reusing the same + // engine. + std::minstd_rand0* engine = nullptr; + + // If pseudo_random is true, the generated numbers will be generated + // deterministically in a pseudo random way unless the values are constrained + // to be e.g. init values as above. If pseudo_random is false, the returned + // values will be generated in a faster way that yields less interesting data, + // e.g. the values may all be just the same value. + // + // TODO(b/79942829): Make interesting argument generation fast enough that + // using pseudo_random does not save any noticeable amount of time so that the + // parameter can be removed. + bool pseudo_random = true; + + // If use_large_range is false, the generated floating point numbers will be + // sampled from a small range of possible values. If use_large_range is true, + // the generated floating point numbers will be sampled from a uniform-log + // distribution of most possible floats, with a small chance to instead be + // sampled from a list of special floating point values (such as 0, inf, + // etc.). + bool use_large_range = false; + + // If treat_gte_as_data_formatting is true, GetTupleElement instructions are + // treated as data formatting operations when tracking parameter constraints, + // allowing constraints to propagate through tuple deconstruction. + bool treat_gte_as_data_formatting = false; + + // If max_bits_of_precision is set to a number, then floating point & integer + // types will be constrained to be represented in that number of bits. Setting + // it to 5 for integers would mean it only creates integers between -32 and + // 32. + std::optional max_bits_of_precision = std::nullopt; + + // If `generate_aligned_ds_indices` is true, the generated indices will be + // aligned to the given alignment. + bool generate_aligned_ds_indices = false; + + // If `get_index_known_zeroes` is set, the generated indices will have the + // given number of zeroes in the given dimension. + GetIndexKnownZeroesFn get_index_known_zeroes = nullptr; +}; + // Generates a vector of arguments containing fake data. The number, shape and // layout of the arguments is appropriate for given HLO module. // @@ -79,80 +126,20 @@ using GetIndexKnownZeroesFn = // (3) Keys of key/value sorts should contain no duplicates. // // These constraints are best-effort only. -// -// If max_bits_of_precision is set to a number, then floating point & integer -// types will be constrained to be represented in that number of bits. Setting -// it to 5 for integers would mean it only creates integers between -32 and 32. -// -// If pseudo_random is true, the generated numbers will be generated -// deterministically in a pseudo random way unless the values are constrated to -// be e.g. init values as above. If pseudo_random is false, the returned values -// will be generated in a faster way that yields less interesting data, e.g. the -// values may all be just the same value. -// -// If use_large_range is false, the generated floating point numbers will be -// sampled from a small range of possible values. If use_large_range is true, -// the generated floating point numbers will be sampled from a uniform-log -// distribution of most possible floats, with a small chance to instead be -// sampled from a list of special floating point values (such as 0, inf, etc.). -// -// TODO(b/79942829): Make interesting argument generation fast enough that using -// pseudo_random does not save any noticeable amount of time so that the -// parameter can be removed. -// -// If `generate_aligned_ds_indices` is true, the generated indices will be -// aligned to the given alignment. If `get_index_known_zeroes` is set, the -// generated indices will have the given number of zeroes in the given -// dimension. -absl::StatusOr> MakeFakeArguments( - const HloModule* module, bool pseudo_random = true, - bool use_large_range = false, bool treat_gte_as_data_formatting = false, - std::optional max_bits_of_precision = std::nullopt, - std::minstd_rand0* engine = nullptr, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); - -// Overload which accepts a random number generator. This enables generation of -// different random values with sequential calls to MakeFakeArguments by reusing -// the same generator. absl::StatusOr> MakeFakeArguments( - const HloModule* module, std::minstd_rand0* engine, - bool use_large_range = false, bool treat_gte_as_data_formatting = false, - std::optional max_bits_of_precision = std::nullopt, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); + const HloModule* module, const FakeArgumentsOptions& options = {}); // Generates a vector of arguments containing fake data using reverse constraint // propagation. The constraint propagator seeds initial constraints based on HLO // op semantics (e.g., `sqrt(x)` implies `x >= 0`) and then propagates these // constraints backward through the graph. This allows generating test inputs // that are more likely to be valid for the graph. -// -// If `use_large_range` is false, the generated floating point numbers will be -// sampled from a small range of possible values. If `use_large_range` is true, -// the generated floating point numbers will be sampled from a uniform-log -// distribution of most possible floats, with a small chance to instead be -// sampled from a list of special floating point values (such as 0, inf, etc.). -// -// If `max_bits_of_precision` is set to a number, then floating point & integer -// types will be constrained to be represented in that number of bits. Setting -// it to 5 for integers would mean it only creates integers between -32 and 32. -// -// If `generate_aligned_ds_indices` is true, the generated indices will be -// aligned to the given alignment. -// -// If `get_index_known_zeroes` is set, the generated indices will have the given -// number of zeroes in the given dimension. absl::StatusOr> MakeDataflowConstrainedArguments( - const HloModule* module, std::minstd_rand0* engine = nullptr, - bool use_large_range = false, - std::optional max_bits_of_precision = std::nullopt, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); + const HloModule* module, const FakeArgumentsOptions& options = {}); // Check that a given module satisfies various constraints before trying to // execute it. -absl::Status VerifyHloModule(HloModule* const module, bool layout_sensitive, +absl::Status VerifyHloModule(HloModule* module, bool layout_sensitive, bool allow_mixed_precision); // Creates a dot op with operands 'lhs' and 'rhs' that contracts dimension 1 of diff --git a/third_party/xla/xla/tests/test_utils_test.cc b/third_party/xla/xla/tests/test_utils_test.cc index 90192dcb520d96..20d7de9c4c62aa 100644 --- a/third_party/xla/xla/tests/test_utils_test.cc +++ b/third_party/xla/xla/tests/test_utils_test.cc @@ -258,11 +258,11 @@ ENTRY cluster_13361217111314620287__.11 { )") .value(); - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/true, - /*treat_gte_as_data_formatting=*/true)); + FakeArgumentsOptions options; + options.use_large_range = true; + options.treat_gte_as_data_formatting = true; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 1); const Shape& indices_shape = args[0].shape().tuple_shapes()[0]; @@ -474,16 +474,10 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s32[]) -> s8[131072,2048] { return std::nullopt; }; - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + FakeArgumentsOptions options; + options.get_index_known_zeroes = index_known_zeroes_fn; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 2); int32_t index = args[1].Get({}); @@ -515,16 +509,10 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) return std::nullopt; }; - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + FakeArgumentsOptions options; + options.get_index_known_zeroes = index_known_zeroes_fn; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 3); int32_t index = args[2].Get({}); @@ -533,12 +521,7 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) TF_ASSERT_OK_AND_ASSIGN( std::vector args2, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + MakeDataflowConstrainedArguments(module.get(), options)); ASSERT_EQ(args2.size(), 3); int32_t index2 = args2[2].Get({}); EXPECT_EQ(index2 & index_known_bits_zero, 0); @@ -559,11 +542,8 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false)); + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeDataflowConstrainedArguments(module.get())); ASSERT_EQ(args.size(), 1); args[0].EachCell([](absl::Span indices, float value) { EXPECT_GT(value, 0.0f); @@ -580,11 +560,8 @@ ENTRY main { } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false)); + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeDataflowConstrainedArguments(module.get())); ASSERT_EQ(args.size(), 1); args[0].EachCell([](absl::Span indices, int32_t value) { EXPECT_GE(value, 1); }); diff --git a/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc b/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc index 7028adf7389728..78d82ddb469450 100644 --- a/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc +++ b/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc @@ -194,8 +194,10 @@ MiscompareChecker::MiscompareChecker(HloModule* module, // Generate input data and store the data for all the execution. std::minstd_rand0 rng_engine; if (input_data.empty()) { + FakeArgumentsOptions options; + options.engine = &rng_engine; absl::StatusOr> input_status = - MakeFakeArguments(module, &rng_engine); + MakeFakeArguments(module, options); CHECK(input_status.ok()); input_data_ = std::move(input_status).value(); } else { 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 ecd7b78a58a240..1fd28ce0b54187 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 @@ -510,23 +510,17 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s32[]) -> s8[131072,2048] { } )")); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - [](const HloInstruction* use, - int64_t sliced_dim) -> std::optional { - if (use->opcode() == HloOpcode::kDynamicSlice && - sliced_dim == 0) { - return 131071; - } - return std::nullopt; - })); + FakeArgumentsOptions options; + options.get_index_known_zeroes = + [](const HloInstruction* use, + int64_t sliced_dim) -> std::optional { + if (use->opcode() == HloOpcode::kDynamicSlice && sliced_dim == 0) { + return 131071; + } + return std::nullopt; + }; + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 2); int32_t index = args[1].Get({}); @@ -548,23 +542,17 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) } )")); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - [](const HloInstruction* use, - int64_t sliced_dim) -> std::optional { - if (use->opcode() == HloOpcode::kDynamicUpdateSlice && - sliced_dim == 0) { - return 131071; - } - return std::nullopt; - })); + FakeArgumentsOptions options_dus; + options_dus.get_index_known_zeroes = + [](const HloInstruction* use, + int64_t sliced_dim) -> std::optional { + if (use->opcode() == HloOpcode::kDynamicUpdateSlice && sliced_dim == 0) { + return 131071; + } + return std::nullopt; + }; + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options_dus)); ASSERT_EQ(args.size(), 3); int32_t index = args[2].Get({}); diff --git a/third_party/xla/xla/tools/matmul_perf_table_gen.cc b/third_party/xla/xla/tools/matmul_perf_table_gen.cc index ac16e18c9df78b..530c137195fae0 100644 --- a/third_party/xla/xla/tools/matmul_perf_table_gen.cc +++ b/third_party/xla/xla/tools/matmul_perf_table_gen.cc @@ -402,23 +402,28 @@ absl::Duration MatmulPerfTableGen::Profile(std::unique_ptr module) { // Flip flop between arguments to prevent caching. std::minstd_rand0 engine; - std::vector args_small = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/false) - .value(); - std::vector args_large = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/true) - .value(); + FakeArgumentsOptions small_options; + small_options.engine = &engine; + small_options.use_large_range = false; + auto args_small = MakeFakeArguments(module.get(), small_options); + CHECK_OK(args_small); + + FakeArgumentsOptions large_options; + large_options.engine = &engine; + large_options.use_large_range = true; + auto args_large = MakeFakeArguments(module.get(), large_options); + CHECK_OK(args_large); std::unique_ptr compiled = Compile(std::move(module)); // First run to warm up stuff. - CHECK_OK(runner_.ExecuteWithExecutable(compiled.get(), args_small).status()); + CHECK_OK(runner_.ExecuteWithExecutable(compiled.get(), *args_small).status()); // Trace `kNumProfilingRuns` times to get decent measurement. std::unique_ptr tracer = HloOpProfiler::GetKernelTracer(); for (int i = 0; i < kNumProfilingRuns; i++) { - Measure(runner_, compiled.get(), args_small, args_large); + Measure(runner_, compiled.get(), *args_small, *args_large); } return absl::Nanoseconds(std::move(*tracer).getMedianKernelTimeNs()); diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc index 9da2396d55903f..8d1ab3ec5e91aa 100644 --- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc +++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc @@ -1004,23 +1004,18 @@ CreateArgumentsOnDevice(PjRtClient& client, } } } else { + FakeArgumentsOptions options; + options.engine = engine; + options.pseudo_random = kUseRandomInputs; if (flatten_arguments) { - ABSL_ASSIGN_OR_RETURN( - LiteralVec tupled_argument_literals, - MakeFakeArguments(my_hlo_module, kUseRandomInputs, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, engine)); + ABSL_ASSIGN_OR_RETURN(LiteralVec tupled_argument_literals, + MakeFakeArguments(my_hlo_module, options)); CHECK_EQ(tupled_argument_literals.size(), 1); CHECK(tupled_argument_literals.front().shape().IsTuple()); argument_literals = tupled_argument_literals.front().DecomposeTuple(); } else { - ABSL_ASSIGN_OR_RETURN( - argument_literals, - MakeFakeArguments(my_hlo_module, kUseRandomInputs, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, engine)); + ABSL_ASSIGN_OR_RETURN(argument_literals, + MakeFakeArguments(my_hlo_module, options)); } if (kUseSharedInputs) { break; diff --git a/third_party/xla/xla/tools/run_hlo_module.cc b/third_party/xla/xla/tools/run_hlo_module.cc index f219c1cbc51884..ffca6b41055396 100644 --- a/third_party/xla/xla/tools/run_hlo_module.cc +++ b/third_party/xla/xla/tools/run_hlo_module.cc @@ -242,12 +242,15 @@ absl::Status RunAndCompareInternal( .status()); } - ABSL_ASSIGN_OR_RETURN(auto args, - copy_result_on_failure( - MakeFakeArguments(test_module.get(), engine, - options.use_large_float_range, - options.treat_gte_as_data_formatting), - ModuleResult::kOtherError, test_run_result)); + FakeArgumentsOptions fake_arguments_options; + fake_arguments_options.engine = engine; + fake_arguments_options.use_large_range = options.use_large_float_range; + fake_arguments_options.treat_gte_as_data_formatting = + options.treat_gte_as_data_formatting; + ABSL_ASSIGN_OR_RETURN(auto args, copy_result_on_failure( + MakeFakeArguments(test_module.get(), + fake_arguments_options), + ModuleResult::kOtherError, test_run_result)); // Use provided input literals as arguments, if any. if (iteration_literals_proto != nullptr && iteration_literals_proto->arguments_size() != 0) {